Fix tests for Expression.subs

[linpy.git] / pypol / linexprs.py

diff --git a/pypol/linexprs.py b/pypol/linexprs.py
index ed68493..b330045 100644 (file)
--- a/pypol/linexprs.py
+++ b/pypol/linexprs.py
@@ -249,6 +249,29 @@ class Expression:
                 return left / right
         raise SyntaxError('invalid syntax')
 
+    def subs(self, symbol, expression=None):
+        if expression is None:
+            if isinstance(symbol, dict):
+                symbol = symbol.items()
+            substitutions = symbol
+        else:
+            substitutions = [(symbol, expression)]
+        result = self
+        for symbol, expression in substitutions:
+            symbol = symbolname(symbol)
+            result = result._subs(symbol, expression)
+        return result
+
+    def _subs(self, symbol, expression):
+        coefficients = {name: coefficient
+            for name, coefficient in self.coefficients()
+            if name != symbol}
+        constant = self.constant
+        coefficient = self.coefficient(symbol)
+        result = Expression(coefficients, self.constant)
+        result += coefficient * expression
+        return result
+
     _RE_NUM_VAR = re.compile(r'(\d+|\))\s*([^\W\d_]\w*|\()')
 
     @classmethod