From 41c146a6f3527d797c74c189966621ba27f0bfca Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Fri, 15 May 2020 19:09:44 +0200
Subject: [PATCH] Implement magic method return values, fixes #1577

---
 jedi/inference/syntax_tree.py             | 43 +++++++++++++++++++----
 test/completion/functions.py              |  2 +-
 test/completion/positional_only_params.py | 10 ++++--
 test/completion/precedence.py             | 17 ++++++++-
 4 files changed, 61 insertions(+), 11 deletions(-)

diff --git a/jedi/inference/syntax_tree.py b/jedi/inference/syntax_tree.py
index 7983bef3..015b1b92 100644
--- a/jedi/inference/syntax_tree.py
+++ b/jedi/inference/syntax_tree.py
@@ -31,6 +31,25 @@ from jedi.inference.context import CompForContext
 from jedi.inference.value.decorator import Decoratee
 from jedi.plugins import plugin_manager
 
+operator_to_magic_method = {
+    '+': '__add__',
+    '-': '__sub__',
+    '*': '__mul__',
+    '/': '__div__',
+    '//': '__floordiv__',
+    '%': '__mod__',
+    '**': '__pow__',
+    '<<': '__lshift__',
+    '>>': '__rshift__',
+    '&': '__and__',
+    '|': '__or__',
+    '^': '__xor__',
+}
+
+reverse_operator_to_magic_method = {
+    k: '__r' + v[2:] for k, v in operator_to_magic_method.items()
+}
+
 
 def _limit_value_infers(func):
     """
@@ -538,12 +557,8 @@ def _is_annotation_name(name):
     return False
 
 
-def _is_tuple(value):
-    return isinstance(value, iterable.Sequence) and value.array_type == 'tuple'
-
-
-def _is_list(value):
-    return isinstance(value, iterable.Sequence) and value.array_type == 'list'
+def _is_list_or_tuple(value):
+    return value.array_type in ('tuple', 'list')
 
 
 def _bool_to_value(inference_state, bool_):
@@ -584,7 +599,7 @@ def _infer_comparison_part(inference_state, context, left, operator, right):
     elif str_operator == '+':
         if l_is_num and r_is_num or is_string(left) and is_string(right):
             return left.execute_operation(right, str_operator)
-        elif _is_tuple(left) and _is_tuple(right) or _is_list(left) and _is_list(right):
+        elif _is_list_or_tuple(left) and _is_list_or_tuple(right):
             return ValueSet([iterable.MergedArray(inference_state, (left, right))])
     elif str_operator == '-':
         if l_is_num and r_is_num:
@@ -637,6 +652,20 @@ def _infer_comparison_part(inference_state, context, left, operator, right):
         analysis.add(context, 'type-error-operation', operator,
                      message % (left, right))
 
+    if left.is_class() or right.is_class():
+        return NO_VALUES
+
+    method_name = operator_to_magic_method[str_operator]
+    magic_methods = left.py__getattribute__(method_name)
+    if not magic_methods:
+        reverse_method_name = reverse_operator_to_magic_method[str_operator]
+        magic_methods = left.py__getattribute__(reverse_method_name)
+
+    if magic_methods:
+        result = magic_methods.execute_with_values(right)
+        if result:
+            return result
+
     result = ValueSet([left, right])
     debug.dbg('Used operator %s resulting in %s', operator, result)
     return result
diff --git a/test/completion/functions.py b/test/completion/functions.py
index 585b692b..9df5bfcf 100644
--- a/test/completion/functions.py
+++ b/test/completion/functions.py
@@ -241,7 +241,7 @@ args_func(*1)[0]
 args_func(*iter([1]))[0]
 
 # different types
-e = args_func(*[1+"", {}])
+e = args_func(*[1 if UNDEFINED else "", {}])
 #? int() str()
 e[0]
 #? dict()
diff --git a/test/completion/positional_only_params.py b/test/completion/positional_only_params.py
index 41261510..3c1108d9 100644
--- a/test/completion/positional_only_params.py
+++ b/test/completion/positional_only_params.py
@@ -5,7 +5,10 @@ def positional_only_call(a, /, b):
     a
     #? int()
     b
-    return a + b
+    if UNDEFINED:
+        return a
+    else:
+        return b
 
 
 #? int() str()
@@ -13,7 +16,10 @@ positional_only_call('', 1)
 
 
 def positional_only_call2(a, /, b=3):
-    return a + b
+    if UNDEFINED:
+        return a
+    else:
+        return b
 
 #? int()
 positional_only_call2(1)
diff --git a/test/completion/precedence.py b/test/completion/precedence.py
index 92c670eb..6c82a2c3 100644
--- a/test/completion/precedence.py
+++ b/test/completion/precedence.py
@@ -56,7 +56,7 @@ a
 (3 ** 3)
 #? int() str()
 (3 ** 'a')
-#? int() str()
+#? int()
 (3 + 'a')
 #? bool()
 (3 == 'a')
@@ -147,3 +147,18 @@ a = foobarbaz + 'hello'
 
 #? int() float()
 {'hello': 1, 'bar': 1.0}[a]
+
+# -----------------
+# stubs
+# -----------------
+
+from datetime import datetime, timedelta
+
+#?
+(datetime - timedelta)
+#? datetime()
+(datetime() - timedelta())
+#? timedelta()
+(timedelta() - datetime())
+#? timedelta()
+(timedelta() - timedelta())