From 0e42df2da732cc016016376ac006bc0cdd73616f Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Wed, 1 May 2019 00:52:02 +0200
Subject: [PATCH] Refactor Jedi so we use stub modules as much as possible

---
 jedi/api/__init__.py                  |  8 ++--
 jedi/evaluate/__init__.py             | 12 ++++--
 jedi/evaluate/compiled/__init__.py    |  8 +++-
 jedi/evaluate/compiled/context.py     |  8 ++++
 jedi/evaluate/context/klass.py        |  7 +---
 jedi/evaluate/context/module.py       |  5 ++-
 jedi/evaluate/gradual/stub_context.py | 52 +++++++++++++++++++++++---
 jedi/evaluate/gradual/typeshed.py     | 54 +++++++++++----------------
 jedi/evaluate/gradual/typing.py       |  3 ++
 jedi/evaluate/gradual/utils.py        |  6 +--
 jedi/evaluate/imports.py              | 36 +++++++++---------
 jedi/plugins/flask.py                 | 12 ++----
 12 files changed, 125 insertions(+), 86 deletions(-)

diff --git a/jedi/api/__init__.py b/jedi/api/__init__.py
index 73d4ae60..1264a678 100644
--- a/jedi/api/__init__.py
+++ b/jedi/api/__init__.py
@@ -38,7 +38,7 @@ from jedi.evaluate.syntax_tree import tree_name_to_contexts
 from jedi.evaluate.context import ModuleContext
 from jedi.evaluate.base_context import ContextSet
 from jedi.evaluate.context.iterable import unpack_tuple_to_dict
-from jedi.evaluate.gradual.typeshed import try_to_merge_with_stub
+#from jedi.evaluate.gradual.typeshed import try_to_merge_with_stub
 from jedi.evaluate.gradual.utils import load_proper_stub_module
 
 # Jedi uses lots and lots of recursion. By setting this a little bit higher, we
@@ -182,9 +182,9 @@ class Script(object):
             code_lines=self._code_lines,
             is_package=is_package,
         )
-        module, = try_to_merge_with_stub(
-            self._evaluator, None, module.string_names, ContextSet([module])
-        )
+        #module, = try_to_merge_with_stub(
+        #    self._evaluator, None, module.string_names, ContextSet([module])
+        #)
         if names[0] not in ('builtins', '__builtin__', 'typing'):
             # These modules are essential for Jedi, so don't overwrite them.
             self._evaluator.module_cache.add(names, ContextSet([module]))
diff --git a/jedi/evaluate/__init__.py b/jedi/evaluate/__init__.py
index bd01da71..62732d37 100644
--- a/jedi/evaluate/__init__.py
+++ b/jedi/evaluate/__init__.py
@@ -86,10 +86,14 @@ from jedi.evaluate.syntax_tree import eval_trailer, eval_expr_stmt, \
     eval_node, check_tuple_assignments
 from jedi.evaluate.gradual.stub_context import with_stub_context_if_possible, \
     stub_to_actual_context_set, goto_with_stubs_if_possible, goto_non_stub, \
-    stubify
+    stubify, load_stubs
 
 
 def _execute(context, arguments):
+    if not context.get_root_context().is_stub():
+        stubs = load_stubs(context)
+        if stubs:
+            return stubs.execute(arguments)
     try:
         func = context.py__call__
     except AttributeError:
@@ -113,6 +117,7 @@ class Evaluator(object):
         self.latest_grammar = parso.load_grammar(version='3.6')
         self.memoize_cache = {}  # for memoize decorators
         self.module_cache = imports.ModuleCache()  # does the job of `sys.modules`.
+        self.stub_module_cache = {}  # Dict[Tuple[str, ...], Optional[ModuleContext]]
         self.compiled_cache = {}  # see `evaluate.compiled.create()`
         self.inferred_element_counts = {}
         self.mixed_cache = {}  # see `evaluate.compiled.mixed._create()`
@@ -141,7 +146,7 @@ class Evaluator(object):
         )
 
     def import_module(self, import_names, parent_module_context=None,
-                      sys_path=None, load_stub=True):
+                      sys_path=None):
         if sys_path is None:
             sys_path = self.get_sys_path()
         try:
@@ -149,8 +154,7 @@ class Evaluator(object):
         except KeyError:
             pass
 
-        context_set = self._import_module(import_names, parent_module_context,
-                                          sys_path, load_stub=load_stub)
+        context_set = self._import_module(import_names, parent_module_context, sys_path)
         self.module_cache.add(import_names, context_set)
         return context_set
 
diff --git a/jedi/evaluate/compiled/__init__.py b/jedi/evaluate/compiled/__init__.py
index a4904c40..e922ae3c 100644
--- a/jedi/evaluate/compiled/__init__.py
+++ b/jedi/evaluate/compiled/__init__.py
@@ -7,8 +7,12 @@ from jedi.evaluate.helpers import execute_evaluated
 
 
 def builtin_from_name(evaluator, string):
-    builtins = evaluator.builtins_module
-    filter_ = next(builtins.get_filters())
+    typing_builtins_module = evaluator.builtins_module
+    if string in ('None', 'True', 'False'):
+        builtins, = typing_builtins_module.non_stub_context_set
+        filter_ = next(builtins.get_filters())
+    else:
+        filter_ = next(typing_builtins_module.get_filters())
     name, = filter_.get(string)
     context, = name.infer()
     return context
diff --git a/jedi/evaluate/compiled/context.py b/jedi/evaluate/compiled/context.py
index cb52700f..bf1238c8 100644
--- a/jedi/evaluate/compiled/context.py
+++ b/jedi/evaluate/compiled/context.py
@@ -82,6 +82,14 @@ class CompiledObject(Context):
     def py__path__(self):
         return self.access_handle.py__path__()
 
+    @property
+    def string_names(self):
+        # For modules
+        return tuple(self.py__name__().split('.'))
+
+    def get_qualified_names(self):
+        return self.string_names
+
     def py__bool__(self):
         return self.access_handle.py__bool__()
 
diff --git a/jedi/evaluate/context/klass.py b/jedi/evaluate/context/klass.py
index cab658c0..bcdf9d56 100644
--- a/jedi/evaluate/context/klass.py
+++ b/jedi/evaluate/context/klass.py
@@ -200,12 +200,7 @@ class ClassMixin(FunctionAndClassMixin):
             from jedi.evaluate.compiled import builtin_from_name
             type_ = builtin_from_name(self.evaluator, u'type')
             if type_ != self:
-                # Return completions of the meta class.
-                yield ClassFilter(
-                    self.evaluator, self, node_context=type_,
-                    origin_scope=origin_scope,
-                    is_instance=is_instance
-                )
+                yield next(type_.get_filters())
 
 
 class ClassContext(use_metaclass(CachedMetaClass, ClassMixin, TreeContext)):
diff --git a/jedi/evaluate/context/module.py b/jedi/evaluate/context/module.py
index b0c85dbd..d432d15a 100644
--- a/jedi/evaluate/context/module.py
+++ b/jedi/evaluate/context/module.py
@@ -53,7 +53,7 @@ class SubModuleDictMixin(object):
             pass
         else:
             for path in method():
-                mods = iter_modules([path])
+                mods = self._iter_modules(path)
                 for module_loader, name, is_pkg in mods:
                     # It's obviously a relative import to the current module.
                     names[name] = SubModuleName(self, name)
@@ -67,6 +67,9 @@ class SubModuleDictMixin(object):
 
         return names
 
+    def _iter_modules(self, path):
+        return iter_modules([path])
+
 
 class ModuleMixin(SubModuleDictMixin):
     def get_filters(self, search_global=False, until_position=None, origin_scope=None):
diff --git a/jedi/evaluate/gradual/stub_context.py b/jedi/evaluate/gradual/stub_context.py
index 3d3ab37f..ef234288 100644
--- a/jedi/evaluate/gradual/stub_context.py
+++ b/jedi/evaluate/gradual/stub_context.py
@@ -1,3 +1,5 @@
+import os
+
 from jedi.cache import memoize_method
 from jedi.parser_utils import get_call_signature_for_any
 from jedi.evaluate.utils import safe_property
@@ -10,7 +12,8 @@ from jedi.evaluate.filters import ParserTreeFilter, \
     NameWrapper, AbstractFilter, TreeNameDefinition
 from jedi.evaluate.compiled.context import CompiledName
 from jedi.evaluate.utils import to_list
-from jedi.evaluate.gradual.typing import TypingModuleFilterWrapper, TypingModuleName
+from jedi.evaluate.gradual.typing import TypingModuleFilterWrapper, \
+    TypingModuleName, AnnotatedClass
 
 
 class _StubContextFilterMixin(object):
@@ -113,13 +116,13 @@ class _StubOnlyContextMixin(object):
     def _get_base_filters(self, filters, search_global=False,
                           until_position=None, origin_scope=None):
         next(filters)  # Ignore the first filter and replace it with our own
-        yield self.get_stub_only_filter(
-            parent_contexts=self.get_stub_contexts(),
-            non_stub_filters=list(self._get_first_non_stub_filters()),
+        stub_only_filters = self._get_stub_only_filters(
             search_global=search_global,
             until_position=until_position,
             origin_scope=origin_scope,
         )
+        for f in stub_only_filters:
+            yield f
 
         for f in filters:
             yield f
@@ -169,6 +172,15 @@ class StubOnlyModuleContext(_StubOnlyContextMixin, ModuleContext):
         for f in self._get_base_filters(filters, search_global, until_position, origin_scope):
             yield f
 
+    def _iter_modules(self, path):
+        dirs = os.listdir(path)
+        for name in dirs:
+            if os.path.isdir(os.path.join(path, name)):
+                yield (None, name, True)
+            if name.endswith('.pyi'):
+                yield (None, name[:-4], True)
+        return []
+
 
 class StubOnlyClass(_StubOnlyContextMixin, ClassMixin, ContextWrapper):
     pass
@@ -209,8 +221,7 @@ class CompiledStubClass(_StubOnlyContextMixin, _CompiledStubContext, ClassMixin)
 
 
 class TypingModuleWrapper(StubOnlyModuleContext):
-    # TODO should use this instead of the isinstance check
-    def get_filterss(self, *args, **kwargs):
+    def get_filters(self, *args, **kwargs):
         filters = super(TypingModuleWrapper, self).get_filters(*args, **kwargs)
         yield TypingModuleFilterWrapper(next(filters))
         for f in filters:
@@ -334,6 +345,35 @@ def stubify(parent_context, context):
         return with_stub_context_if_possible(context)
 
 
+def _load_or_get_stub_module(evaluator, names):
+    return evaluator.stub_module_cache.get(names)
+
+
+def load_stubs(context):
+    root_context = context.get_root_context()
+    stub_module = _load_or_get_stub_module(
+        context.evaluator,
+        root_context.string_names
+    )
+    if stub_module is None:
+        return NO_CONTEXTS
+
+    qualified_names = context.get_qualified_names()
+    if qualified_names is None:
+        return NO_CONTEXTS
+
+    stub_contexts = ContextSet([stub_module])
+    for name in qualified_names:
+        stub_contexts = stub_contexts.py__getattribute__(name)
+
+    if isinstance(context, AnnotatedClass):
+        return ContextSet([
+            context.annotate_other_class(c) if c.is_class() else c
+            for c in stub_contexts
+        ])
+    return stub_contexts
+
+
 class CompiledStubName(NameWrapper):
     def __init__(self, parent_context, compiled_name, stub_name):
         super(CompiledStubName, self).__init__(stub_name)
diff --git a/jedi/evaluate/gradual/typeshed.py b/jedi/evaluate/gradual/typeshed.py
index dd000fc2..4a10f73e 100644
--- a/jedi/evaluate/gradual/typeshed.py
+++ b/jedi/evaluate/gradual/typeshed.py
@@ -6,7 +6,6 @@ from jedi._compatibility import FileNotFoundError
 from jedi.parser_utils import get_cached_code_lines
 from jedi.evaluate.cache import evaluator_function_cache
 from jedi.evaluate.base_context import ContextSet, NO_CONTEXTS
-from jedi.evaluate.context import ModuleContext
 from jedi.evaluate.gradual.stub_context import StubModuleContext, \
     TypingModuleWrapper, StubOnlyModuleContext
 
@@ -89,13 +88,7 @@ def _cache_stub_file_map(version_info):
 
 
 def import_module_decorator(func):
-    def wrapper(evaluator, import_names, parent_module_context, sys_path, load_stub=True):
-        if import_names == ('_sqlite3',):
-            # TODO Maybe find a better solution for this?
-            # The problem is IMO how star imports are priorized and that
-            # there's no clear ordering.
-            return NO_CONTEXTS
-
+    def wrapper(evaluator, import_names, parent_module_context, sys_path):
         if import_names == ('os', 'path'):
             # This is a huge exception, we follow a nested import
             # ``os.path``, because it's a very important one in Python
@@ -103,41 +96,33 @@ def import_module_decorator(func):
             # ``os``.
             if parent_module_context is None:
                 parent_module_context, = evaluator.import_module(('os',))
-            return parent_module_context.py__getattribute__('path')
-
-        from jedi.evaluate.imports import JediImportError
-        try:
-            context_set = func(
+            actual_context_set = parent_module_context.py__getattribute__('path')
+        else:
+            actual_context_set = func(
                 evaluator,
                 import_names,
                 parent_module_context,
                 sys_path,
             )
-        except JediImportError:
-            if import_names == ('typing',):
-                # TODO this is also quite ugly, please refactor.
-                context_set = NO_CONTEXTS
-            else:
-                raise
+        stub = _try_to_load_stub(evaluator, actual_context_set, parent_module_context, import_names)
+        if stub is not None:
+            return ContextSet(stub)
+        return actual_context_set
 
-        if not load_stub:
-            return context_set
-        return try_to_merge_with_stub(evaluator, parent_module_context,
-                                      import_names, context_set)
     return wrapper
 
 
-def try_to_merge_with_stub(evaluator, parent_module_context, import_names, actual_context_set):
+def _try_to_load_stub(evaluator, actual_context_set, parent_module_context, import_names):
     import_name = import_names[-1]
     map_ = None
     if len(import_names) == 1:
         map_ = _cache_stub_file_map(evaluator.grammar.version_info)
-    elif isinstance(parent_module_context, StubModuleContext):
-        if not parent_module_context.stub_context.is_package:
+    elif isinstance(parent_module_context, StubOnlyModuleContext):
+        if not parent_module_context.is_package:
             # Only if it's a package (= a folder) something can be
             # imported.
-            return actual_context_set
-        path = parent_module_context.stub_context.py__path__()
+            return None
+        path = parent_module_context.py__path__()
         map_ = _merge_create_stub_map(path)
 
     if map_ is not None:
@@ -150,10 +135,14 @@ def try_to_merge_with_stub(evaluator, parent_module_context, import_names, actua
                 # TODO maybe empty cache?
                 pass
             else:
-                return create_stub_module(evaluator, actual_context_set,
-                                          stub_module_node, path, import_names)
+                return create_stub_module(
+                    evaluator, actual_context_set, stub_module_node, path,
+                    import_names
+                )
+    evaluator.stub_module_cache[import_names] = None
     # If no stub is found, just return the default.
-    return actual_context_set
+
+    return None
 
 
 def create_stub_module(evaluator, actual_context_set, stub_module_node, path, import_names):
@@ -171,4 +160,5 @@ def create_stub_module(evaluator, actual_context_set, stub_module_node, path, im
         code_lines=get_cached_code_lines(evaluator.latest_grammar, path),
         is_package=file_name == '__init__.pyi',
     )
-    return stub_module_context.get_stub_contexts()
+    evaluator.stub_module_cache[import_names] = stub_module_context
+    return [stub_module_context]
diff --git a/jedi/evaluate/gradual/typing.py b/jedi/evaluate/gradual/typing.py
index 5b1c0187..7939678f 100644
--- a/jedi/evaluate/gradual/typing.py
+++ b/jedi/evaluate/gradual/typing.py
@@ -625,6 +625,9 @@ class AnnotatedClass(AbstractAnnotatedClass):
     def get_given_types(self):
         return list(_iter_over_arguments(self._index_context, self._context_of_index))
 
+    def annotate_other_class(self, cls):
+        return AnnotatedClass(cls, self._index_context, self._context_of_index)
+
 
 class AnnotatedSubClass(AbstractAnnotatedClass):
     def __init__(self, class_context, given_types):
diff --git a/jedi/evaluate/gradual/utils.py b/jedi/evaluate/gradual/utils.py
index 4f3eb795..f64ec046 100644
--- a/jedi/evaluate/gradual/utils.py
+++ b/jedi/evaluate/gradual/utils.py
@@ -1,6 +1,5 @@
 import os
 
-from jedi.evaluate.imports import JediImportError
 from jedi.evaluate.gradual.typeshed import TYPESHED_PATH, create_stub_module
 
 
@@ -20,9 +19,8 @@ def load_proper_stub_module(evaluator, path, import_names, module_node):
             import_names = import_names[:-1]
 
     if import_names is not None:
-        try:
-            actual_context_set = evaluator.import_module(import_names, load_stub=False)
-        except JediImportError as e:
+        actual_context_set = evaluator.import_module(import_names)
+        if not actual_context_set:
             return None
 
         context_set = create_stub_module(
diff --git a/jedi/evaluate/imports.py b/jedi/evaluate/imports.py
index 99760554..4579b596 100644
--- a/jedi/evaluate/imports.py
+++ b/jedi/evaluate/imports.py
@@ -327,16 +327,15 @@ class Importer(object):
 
         context_set = [None]
         for i, name in enumerate(self.import_path):
-            try:
-                context_set = ContextSet.from_sets([
-                    self._evaluator.import_module(
-                        import_names[:i+1],
-                        parent_module_context,
-                        self._sys_path_with_modifications(),
-                    )
-                    for parent_module_context in context_set
-                ])
-            except JediImportError:
+            context_set = ContextSet.from_sets([
+                self._evaluator.import_module(
+                    import_names[:i+1],
+                    parent_module_context,
+                    self._sys_path_with_modifications(),
+                )
+                for parent_module_context in context_set
+            ])
+            if not context_set:
                 message = 'No module named ' + '.'.join(import_names)
                 _add_error(self.module_context, name, message)
                 return NO_CONTEXTS
@@ -424,11 +423,6 @@ class Importer(object):
         return names
 
 
-class JediImportError(Exception):
-    def __init__(self, import_names):
-        self.import_names = import_names
-
-
 @import_module_decorator
 def import_module(evaluator, import_names, parent_module_context, sys_path, load_stub=True):
     """
@@ -436,6 +430,8 @@ def import_module(evaluator, import_names, parent_module_context, sys_path, load
     """
     if import_names[0] in settings.auto_import_modules:
         module = _load_builtin_module(evaluator, import_names, sys_path)
+        if module is None:
+            return NO_CONTEXTS
         return ContextSet([module])
 
     module_name = '.'.join(import_names)
@@ -449,13 +445,13 @@ def import_module(evaluator, import_names, parent_module_context, sys_path, load
             is_global_search=True,
         )
         if is_pkg is None:
-            raise JediImportError(import_names)
+            return NO_CONTEXTS
     else:
         try:
             method = parent_module_context.py__path__
         except AttributeError:
             # The module is not a package.
-            raise JediImportError(import_names)
+            return NO_CONTEXTS
         else:
             paths = method()
             for path in paths:
@@ -472,7 +468,7 @@ def import_module(evaluator, import_names, parent_module_context, sys_path, load
                 if is_pkg is not None:
                     break
             else:
-                raise JediImportError(import_names)
+                return NO_CONTEXTS
 
     if isinstance(file_io_or_ns, ImplicitNSInfo):
         from jedi.evaluate.context.namespace import ImplicitNamespaceContext
@@ -483,6 +479,8 @@ def import_module(evaluator, import_names, parent_module_context, sys_path, load
         )
     elif file_io_or_ns is None:
         module = _load_builtin_module(evaluator, import_names, sys_path)
+        if module is None:
+            return NO_CONTEXTS
     else:
         module = _load_python_module(
             evaluator, file_io_or_ns, sys_path,
@@ -531,7 +529,7 @@ def _load_builtin_module(evaluator, import_names=None, sys_path=None):
     if module is None:
         # The file might raise an ImportError e.g. and therefore not be
         # importable.
-        raise JediImportError(import_names)
+        return None
     return module
 
 
diff --git a/jedi/plugins/flask.py b/jedi/plugins/flask.py
index b082f40f..c4906c4c 100644
--- a/jedi/plugins/flask.py
+++ b/jedi/plugins/flask.py
@@ -1,5 +1,4 @@
 from jedi.plugins.base import BasePlugin
-from jedi.evaluate.imports import JediImportError
 
 
 class FlaskPlugin(BasePlugin):
@@ -8,21 +7,18 @@ class FlaskPlugin(BasePlugin):
         Handle "magic" Flask extension imports:
         ``flask.ext.foo`` is really ``flask_foo`` or ``flaskext.foo``.
         """
-        def wrapper(evaluator, import_names, module_context, sys_path, load_stub):
+        def wrapper(evaluator, import_names, module_context, sys_path):
             if len(import_names) == 3 and import_names[:2] == ('flask', 'ext'):
                 # New style.
                 ipath = (u'flask_' + import_names[2]),
-                try:
-                    return callback(evaluator, ipath, None, sys_path, load_stub)
-                except JediImportError:
+                context_set = callback(evaluator, ipath, None, sys_path)
+                if not context_set:
                     context_set = callback(evaluator, (u'flaskext',), None, sys_path)
-                    # If context_set has no content a JediImportError is raised
-                    # which should be caught anyway by the caller.
                     return callback(
                         evaluator,
                         (u'flaskext', import_names[2]),
                         next(iter(context_set)),
                         sys_path
                     )
-            return callback(evaluator, import_names, module_context, sys_path, load_stub)
+            return callback(evaluator, import_names, module_context, sys_path)
         return wrapper