The inference engine handles binary operations (add, mul etc.) between instances.

4 files changed, 104 insertions, 24 deletions

diff --git a/ChangeLog b/ChangeLog
index 55384a0..ce9bf3e 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -19,6 +19,9 @@ Change log for the astroid package (used to be astng)
     * Inference of the functional form of the enums takes into
       consideration the various inputs that enums accepts.
 
+    * The inference engine handles binary operations (add, mul etc.)
+      between instances.
+
 2014-07-25  --  1.2.0
 
     * Function nodes can detect decorator call chain and see if they are
diff --git a/inference.py b/inference.py
index c36f695..4186307 100644
--- a/inference.py
+++ b/inference.py
@@ -30,7 +30,9 @@ from astroid.exceptions import (AstroidError, InferenceError, NoDefault,
 from astroid.bases import (YES, Instance, InferenceContext,
                            _infer_stmts, copy_context, path_wrapper,
                            raise_if_nothing_infered)
-from astroid.protocols import _arguments_infer_argname
+from astroid.protocols import (
+    _arguments_infer_argname,
+    BIN_OP_METHOD, UNARY_OP_METHOD)
 
 MANAGER = AstroidManager()
 
@@ -292,13 +294,6 @@ def infer_subscript(self, context=None):
 nodes.Subscript._infer = path_wrapper(infer_subscript)
 nodes.Subscript.infer_lhs = raise_if_nothing_infered(infer_subscript)
 
-
-UNARY_OP_METHOD = {'+': '__pos__',
-                   '-': '__neg__',
-                   '~': '__invert__',
-                   'not': None, # XXX not '__nonzero__'
-                  }
-
 def infer_unaryop(self, context=None):
     for operand in self.operand.infer(context):
         try:
@@ -321,21 +316,6 @@ def infer_unaryop(self, context=None):
                     yield YES
 nodes.UnaryOp._infer = path_wrapper(infer_unaryop)
 
-
-BIN_OP_METHOD = {'+':  '__add__',
-                 '-':  '__sub__',
-                 '/':  '__div__',
-                 '//': '__floordiv__',
-                 '*':  '__mul__',
-                 '**': '__power__',
-                 '%':  '__mod__',
-                 '&':  '__and__',
-                 '|':  '__or__',
-                 '^':  '__xor__',
-                 '<<': '__lshift__',
-                 '>>': '__rshift__',
-                }
-
 def _infer_binop(operator, operand1, operand2, context, failures=None):
     if operand1 is YES:
         yield operand1
diff --git a/protocols.py b/protocols.py
index d621ffb..e7703a0 100644
--- a/protocols.py
+++ b/protocols.py
@@ -21,13 +21,33 @@ where it makes sense.
 
 __doctype__ = "restructuredtext en"
 
-from astroid.exceptions import InferenceError, NoDefault
+from astroid.exceptions import InferenceError, NoDefault, NotFoundError
 from astroid.node_classes import unpack_infer
 from astroid.bases import copy_context, \
      raise_if_nothing_infered, yes_if_nothing_infered, Instance, YES
 from astroid.nodes import const_factory
 from astroid import nodes
 
+BIN_OP_METHOD = {'+':  '__add__',
+                 '-':  '__sub__',
+                 '/':  '__div__',
+                 '//': '__floordiv__',
+                 '*':  '__mul__',
+                 '**': '__power__',
+                 '%':  '__mod__',
+                 '&':  '__and__',
+                 '|':  '__or__',
+                 '^':  '__xor__',
+                 '<<': '__lshift__',
+                 '>>': '__rshift__',
+                }
+
+UNARY_OP_METHOD = {'+': '__pos__',
+                   '-': '__neg__',
+                   '~': '__invert__',
+                   'not': None, # XXX not '__nonzero__'
+                  }
+
 # unary operations ############################################################
 
 def tl_infer_unary_op(self, operator):
@@ -133,6 +153,25 @@ def dict_infer_binary_op(self, operator, other, context):
         # XXX else log TypeError
 nodes.Dict.infer_binary_op = yes_if_nothing_infered(dict_infer_binary_op)
 
+def instance_infer_binary_op(self, operator, other, context):
+    try:
+        methods = self.getattr(BIN_OP_METHOD[operator])
+    except (NotFoundError, KeyError):
+        # Unknown operator
+        yield YES
+    else:
+        for method in methods:
+            if not isinstance(method, nodes.Function):
+                continue
+            for result in method.infer_call_result(self, context):
+                if result is not YES:
+                    yield result
+            # We are interested only in the first infered method,
+            # don't go looking in the rest of the methods of the ancestors.
+            break
+
+Instance.infer_binary_op = yes_if_nothing_infered(instance_infer_binary_op)
+
 
 # assignment ##################################################################
 
diff --git a/test/unittest_inference.py b/test/unittest_inference.py
index 06f5abb..ea49e50 100644
--- a/test/unittest_inference.py
+++ b/test/unittest_inference.py
@@ -1257,5 +1257,63 @@ def test(*args, **kwargs):
         infered = func.infered()[0]
         self.assertIsInstance(infered, UnboundMethod)
 
+    def test_instance_binary_operations(self):
+        code = dedent("""
+        class A(object):
+            def __mul__(self, other):
+                return 42
+        a = A()
+        b = A()
+        sub = a - b
+        mul = a * b
+        """)
+        astroid = builder.string_build(code, __name__, __file__)
+        sub = astroid['sub'].infered()[0]
+        mul = astroid['mul'].infered()[0]
+        self.assertIs(sub, YES)
+        self.assertIsInstance(mul, nodes.Const)
+        self.assertEqual(mul.value, 42)
+
+    def test_instance_binary_operations_parent(self):
+        code = dedent("""
+        class A(object):
+            def __mul__(self, other):
+                return 42
+        class B(A):
+            pass
+        a = B()
+        b = B()
+        sub = a - b
+        mul = a * b
+        """)
+        astroid = builder.string_build(code, __name__, __file__)
+        sub = astroid['sub'].infered()[0]
+        mul = astroid['mul'].infered()[0]
+        self.assertIs(sub, YES)
+        self.assertIsInstance(mul, nodes.Const)
+        self.assertEqual(mul.value, 42)
+
+    def test_instance_binary_operations_multiple_methods(self):
+        code = dedent("""
+        class A(object):
+            def __mul__(self, other):
+                return 42
+        class B(A):
+            def __mul__(self, other):
+                return [42]
+        a = B()
+        b = B()
+        sub = a - b
+        mul = a * b
+        """)
+        astroid = builder.string_build(code, __name__, __file__)
+        sub = astroid['sub'].infered()[0]
+        mul = astroid['mul'].infered()[0]
+        self.assertIs(sub, YES)
+        self.assertIsInstance(mul, nodes.List)
+        self.assertIsInstance(mul.elts[0], nodes.Const)
+        self.assertEqual(mul.elts[0].value, 42)
+
+
 if __name__ == '__main__':
     unittest_main()