Fix LinExpr.subs() implementation

[linpy.git] / linpy / linexprs.py

diff --git a/linpy/linexprs.py b/linpy/linexprs.py
index ab5d344..b5864e1 100644 (file)
--- a/linpy/linexprs.py
+++ b/linpy/linexprs.py
@@ -178,8 +178,7 @@ class LinExpr:
         Iterate over the coefficient values in the expression, and the constant
         term.
         """
-        for coefficient in self._coefficients.values():
-            yield coefficient
+        yield from self._coefficients.values()
         yield self._constant
 
     def __bool__(self):
@@ -248,9 +247,10 @@ class LinExpr:
         """
         Test whether two linear expressions are equal.
         """
-        return isinstance(other, LinExpr) and \
-            self._coefficients == other._coefficients and \
-            self._constant == other._constant
+        if isinstance(other, LinExpr):
+            return self._coefficients == other._coefficients and \
+                self._constant == other._constant
+        return NotImplemented
 
     def __le__(self, other):
         from .polyhedra import Le
@@ -273,9 +273,9 @@ class LinExpr:
         Return the expression multiplied by its lowest common denominator to
         make all values integer.
         """
-        lcm = functools.reduce(lambda a, b: a*b // gcd(a, b),
+        lcd = functools.reduce(lambda a, b: a*b // gcd(a, b),
             [value.denominator for value in self.values()])
-        return self * lcm
+        return self * lcd
 
     def subs(self, symbol, expression=None):
         """
@@ -294,21 +294,16 @@ class LinExpr:
         2*x + y + 1
         """
         if expression is None:
-            if isinstance(symbol, Mapping):
-                symbol = symbol.items()
-            substitutions = symbol
+            substitutions = dict(symbol)
         else:
-            substitutions = [(symbol, expression)]
-        result = self
-        for symbol, expression in substitutions:
+            substitutions = {symbol: expression}
+        for symbol in substitutions:
             if not isinstance(symbol, Symbol):
                 raise TypeError('symbols must be Symbol instances')
-            coefficients = [(othersymbol, coefficient)
-                for othersymbol, coefficient in result._coefficients.items()
-                if othersymbol != symbol]
-            coefficient = result._coefficients.get(symbol, 0)
-            constant = result._constant
-            result = LinExpr(coefficients, constant) + coefficient*expression
+        result = self._constant
+        for symbol, coefficient in self._coefficients.items():
+            expression = substitutions.get(symbol, symbol)
+            result += coefficient * expression
         return result
 
     @classmethod
@@ -498,7 +493,9 @@ class Symbol(LinExpr):
         return True
 
     def __eq__(self, other):
-        return self.sortkey() == other.sortkey()
+        if isinstance(other, Symbol):
+            return self.sortkey() == other.sortkey()
+        return NotImplemented
 
     def asdummy(self):
         """