From 3ffe8475b8b4780c60b660afab165dda2ceb9e2f Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Sun, 12 Jul 2020 22:20:06 +0200
Subject: [PATCH] Make sure the interpreter completions work better in Jupyter
 Notebook, fixes #1628

---
 jedi/api/interpreter.py           | 35 +++++++++++++++++++++++++++----
 test/test_api/test_interpreter.py | 11 ++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/jedi/api/interpreter.py b/jedi/api/interpreter.py
index 101c696d..0f944ccf 100644
--- a/jedi/api/interpreter.py
+++ b/jedi/api/interpreter.py
@@ -3,7 +3,9 @@ TODO Some parts of this module are still not well documented.
 """
 
 from jedi.inference import compiled
+from jedi.inference.base_value import ValueSet
 from jedi.inference.filters import ParserTreeFilter, MergedFilter
+from jedi.inference.names import TreeNameDefinition
 from jedi.inference.compiled import mixed
 from jedi.inference.compiled.access import create_access_path
 from jedi.inference.context import ModuleContext
@@ -20,10 +22,36 @@ class NamespaceObject(object):
         self.__dict__ = dct
 
 
+class MixedTreeName(TreeNameDefinition):
+    def infer(self):
+        """
+        In IPython notebook it is typical that some parts of the code that is
+        provided was already executed. In that case if something is not properly
+        inferred, it should still infer from the variables it already knows.
+        """
+        inferred = super(MixedTreeName, self).infer()
+        if not inferred:
+            for compiled_value in self.parent_context.compiled_values:
+                for f in compiled_value.get_filters():
+                    values = ValueSet.from_sets(
+                        n.infer() for n in f.get(self.string_name)
+                    )
+                    if values:
+                        return values
+        return inferred
+
+
+class MixedParserTreeFilter(ParserTreeFilter):
+    name_class = MixedTreeName
+
+
 class MixedModuleContext(ModuleContext):
     def __init__(self, tree_module_value, namespaces):
         super(MixedModuleContext, self).__init__(tree_module_value)
-        self._namespace_objects = [NamespaceObject(n) for n in namespaces]
+        self.compiled_values = [
+            _create(self.inference_state, NamespaceObject(n))
+            for n in namespaces
+        ]
 
     def _get_mixed_object(self, compiled_value):
         return mixed.MixedObject(
@@ -33,7 +61,7 @@ class MixedModuleContext(ModuleContext):
 
     def get_filters(self, until_position=None, origin_scope=None):
         yield MergedFilter(
-            ParserTreeFilter(
+            MixedParserTreeFilter(
                 parent_context=self,
                 until_position=until_position,
                 origin_scope=origin_scope
@@ -41,8 +69,7 @@ class MixedModuleContext(ModuleContext):
             self.get_global_filter(),
         )
 
-        for namespace_obj in self._namespace_objects:
-            compiled_value = _create(self.inference_state, namespace_obj)
+        for compiled_value in self.compiled_values:
             mixed_object = self._get_mixed_object(compiled_value)
             for filter in mixed_object.get_filters(until_position, origin_scope):
                 yield filter
diff --git a/test/test_api/test_interpreter.py b/test/test_api/test_interpreter.py
index e6782e7f..1bd06405 100644
--- a/test/test_api/test_interpreter.py
+++ b/test/test_api/test_interpreter.py
@@ -683,6 +683,17 @@ def test_string_annotation(annotations, result, code):
     assert [d.name for d in defs] == result
 
 
+def test_name_not_inferred_properly():
+    """
+    In IPython notebook it is typical that some parts of the code that is
+    provided was already executed. In that case if something is not properly
+    inferred, it should still infer from the variables it already knows.
+    """
+    x = 1
+    d, = jedi.Interpreter('x = UNDEFINED; x', [locals()]).infer()
+    assert d.name == 'int'
+
+
 def test_variable_reuse():
     x = 1
     d, = jedi.Interpreter('y = x\ny', [locals()]).infer()