[PATCH] Added optional boolean argument to coeff to control whether .expand() is called

[Luke]

Attached is a patch that allows for an optional boolean argument to
coeff which will control whether the expression will be expanded or
not.  The default behavior is to expand everything.  The reason this
is useful is if you have an expression which is linear in a term that
isn't nested, so you don't need to expand everything out and then
collect terms.  Especially when you have coefficients which may be
long and complicated or are already in a nice compact form, you don't
want to expand them out.

~Luke

--~--~---------~--~----~------------~-------~--~----~
You received this message because you are subscribed to the Google Groups 
"sympy-patches" group.
To post to this group, send email to sympy-patches@googlegroups.com
To unsubscribe from this group, send email to 
sympy-patches+unsubscr...@googlegroups.com
For more options, visit this group at 
http://groups.google.com/group/sympy-patches?hl=en
-~----------~----~----~----~------~----~------~--~---


From 96ba1915506155a6036606e875d01e33bc867b43 Mon Sep 17 00:00:00 2001
From: Luke Peterson <hazelnu...@gmail.com>
Date: Sat, 8 Aug 2009 09:06:33 -0700
Subject: [PATCH] Added optional boolean argument to coeff to control whether .expand() is called

---
 sympy/core/basic.py            |    6 ++++--
 sympy/core/tests/test_basic.py |    9 +++++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/sympy/core/basic.py b/sympy/core/basic.py
index 85cec0b..ac1e256 100644
--- a/sympy/core/basic.py
+++ b/sympy/core/basic.py
@@ -1589,7 +1589,7 @@ class Basic(AssumeMeths):
                 else:
                     return self
 
-    def coeff(self, x):
+    def coeff(self, x, expand=True):
         """
         Returns the coefficient of the term "x" or None if there is no "x".
 
@@ -1616,7 +1616,9 @@ class Basic(AssumeMeths):
                 return None
         if x.is_Integer:
             return
-        self = self.expand() # collect expects it's arguments in expanded form
+
+        if expand:
+            self = self.expand() # collect expects it's arguments in expanded form
         result = collect(self, x, evaluate=False)
         if x in result:
             return result[x]
diff --git a/sympy/core/tests/test_basic.py b/sympy/core/tests/test_basic.py
index 238f144..addceee 100644
--- a/sympy/core/tests/test_basic.py
+++ b/sympy/core/tests/test_basic.py
@@ -621,6 +621,15 @@ def test_coeff2_0():
 
     assert g.coeff(psi(r).diff(r, 2)) == 1
 
+def test_coeff_expand():
+    x, y, z = symbols('x y z')
+    expr = z*(x+y)**2
+    expr2 = z*(x+y)**2 + z*(2*x + 2*y)**2
+    assert expr.coeff(z) == 2*x*y + x**2 + y**2
+    assert expr.coeff(z, expand=False) == (x+y)**2
+    assert expr2.coeff(z) == 10*x*y + 5*x**2 + 5*y**2
+    assert expr2.coeff(z, expand=False) == (x+y)**2 + (2*x + 2*y)**2
+
 def test_integrate():
     assert (log(x)).integrate((x, 0, 1)) == -1
     assert sin(x).integrate(x) == -cos(x)
-- 
1.6.0.4