Fix stub searching for nested modules

jedi/evaluate/context/module.py

@@ -88,11 +88,13 @@ class ModuleContext(TreeContext):
     @property
     def _string_name(self):
         """ This is used for the goto functions. """
+        # TODO It's ugly that we even use this, the name is usually well known
+        # ahead so just pass it when create a ModuleContext.
         if self._path is None:
             return ''  # no path -> empty name
         else:
             sep = (re.escape(os.path.sep),) * 2
-            r = re.search(r'([^%s]*?)(%s__init__)?(\.py|\.so)?$' % sep, self._path)
+            r = re.search(r'([^%s]*?)(%s__init__)?(\.pyi?|\.so)?$' % sep, self._path)
             # Remove PEP 3149 names
             return re.sub(r'\.[a-z]+-\d{2}[mud]{0,3}$', '', r.group(1))
 
@@ -106,7 +108,7 @@ class ModuleContext(TreeContext):
         :return: The path to the directory of a package. None in case it's not
                  a package.
         """
-        for suffix in all_suffixes():
+        for suffix in all_suffixes() + ['.pyi']:
             ending = '__init__' + suffix
             py__file__ = self.py__file__()
             if py__file__ is not None and py__file__.endswith(ending):
@@ -139,7 +141,7 @@ class ModuleContext(TreeContext):
     def _py__path__(self):
         search_path = self.evaluator.get_sys_path()
         init_path = self.py__file__()
-        if os.path.basename(init_path) == '__init__.py':
+        if os.path.basename(init_path) in ('__init__.py', '__init__.pyi'):
             with open(init_path, 'rb') as f:
                 content = python_bytes_to_unicode(f.read(), errors='replace')
                 # these are strings that need to be used for namespace packages,

jedi/evaluate/imports.py

@@ -293,10 +293,10 @@ class Importer(object):
                     self._evaluator.import_module(
                         self._evaluator,
                         import_names[:i+1],
-                        module_context,
+                        parent_module_context,
                         self.sys_path_with_modifications(),
                     )
-                    for module_context in context_set
+                    for parent_module_context in context_set
                 ])
             except JediImportError:
                 _add_error(self.module_context, name)
@@ -397,7 +397,7 @@ class JediImportError(Exception):
         self.import_names = import_names
 
 
-def import_module(evaluator, import_names, module_context, sys_path):
+def import_module(evaluator, import_names, parent_module_context, sys_path):
     """
     This method is very similar to importlib's `_gcd_import`.
     """
@@ -415,7 +415,7 @@ def import_module(evaluator, import_names, module_context, sys_path):
     except KeyError:
         pass
 
-    if module_context is None:
+    if parent_module_context is None:
         debug.dbg('global search_module %s', import_names[-1])
         # Override the sys.path. It works only good that way.
         # Injecting the path directly into `find_module` did not work.
@@ -429,7 +429,7 @@ def import_module(evaluator, import_names, module_context, sys_path):
             raise JediImportError(import_names)
     else:
         try:
-            method = module_context.py__path__
+            method = parent_module_context.py__path__
         except AttributeError:
             # The module is not a package.
             raise JediImportError(import_names)

jedi/plugins/stdlib.py

@@ -78,14 +78,14 @@ class StdlibPlugin(BasePlugin):
         return wrapper
 
     def import_module(self, callback):
-        def wrapper(evaluator, import_names, module_context, sys_path):
+        def wrapper(evaluator, import_names, parent_module_context, sys_path):
             # This is a huge exception, we follow a nested import
             # ``os.path``, because it's a very important one in Python
             # that is being achieved by messing with ``sys.modules`` in
             # ``os``.
             if import_names == ('os', 'path'):
-                return module_context.py__getattribute__('path')
-            return callback(evaluator, import_names, module_context, sys_path)
+                return parent_module_context.py__getattribute__('path')
+            return callback(evaluator, import_names, parent_module_context, sys_path)
 
 
         return wrapper

jedi/plugins/typeshed.py

@@ -4,7 +4,7 @@ from pkg_resources import resource_filename
 
 from jedi._compatibility import FileNotFoundError
 from jedi.plugins.base import BasePlugin
-from jedi.evaluate.cache import evaluator_as_method_param_cache
+from jedi.evaluate.cache import evaluator_function_cache
 from jedi.evaluate.base_context import Context, ContextSet, NO_CONTEXTS
 from jedi.evaluate.context import ModuleContext
 
@@ -12,6 +12,13 @@ from jedi.evaluate.context import ModuleContext
 _TYPESHED_PATH = resource_filename('jedi', os.path.join('third_party', 'typeshed'))
 
 
+def _merge_create_stub_map(directories):
+    map_ = {}
+    for directory in directories:
+        map_.update(_create_stub_map(directory))
+    return map_
+
+
 def _create_stub_map(directory):
     """
     Create a mapping of an importable name in Python to a stub file.
@@ -53,6 +60,11 @@ def _get_typeshed_directories(version_info):
             yield os.path.join(base, check_version)
 
 
+@evaluator_function_cache()
+def _load_stub(evaluator, path):
+    return evaluator.parse(path=path, cache=True)
+
+
 class TypeshedPlugin(BasePlugin):
     _version_cache = {}
 
@@ -68,48 +80,53 @@ class TypeshedPlugin(BasePlugin):
         except KeyError:
             pass
 
-        self._version_cache[version] = file_set = {}
-        for dir_ in _get_typeshed_directories(version_info):
-            file_set.update(_create_stub_map(dir_))
-
+        self._version_cache[version] = file_set = \
+            _merge_create_stub_map(_get_typeshed_directories(version_info))
         return file_set
 
-    @evaluator_as_method_param_cache()
-    def _load_stub(self, evaluator, path):
-        return evaluator.parse(path=path, cache=True)
-
     def import_module(self, callback):
-        def wrapper(evaluator, import_names, module_context, sys_path):
+        def wrapper(evaluator, import_names, parent_module_context, sys_path):
             # This is a huge exception, we follow a nested import
             # ``os.path``, because it's a very important one in Python
             # that is being achieved by messing with ``sys.modules`` in
             # ``os``.
-            mapped = self._cache_stub_file_map(evaluator.grammar.version_info)
-            context_set = callback(evaluator, import_names, module_context, sys_path)
-            if len(import_names) == 1 and import_names[0] != 'typing':
-                path = mapped.get(import_names[0])
+            def _find_and_load_stub_module(stub_map):
+                path = stub_map.get(import_name)
                 if path is not None:
                     try:
-                        stub_module = self._load_stub(evaluator, path)
+                        stub_module = _load_stub(evaluator, path)
                     except FileNotFoundError:
                         # The file has since been removed after looking for it.
                         # TODO maybe empty cache?
-                        pass
+                        return None
                     else:
                         return ContextSet.from_iterable(
-                            StubProxy(
-                                context.parent_context,
+                            ModuleStubProxy(
+                                parent_module_context,
                                 context,
-                                ModuleContext(evaluator, stub_module, path, code_lines=[])
+                                ModuleContext(evaluator, stub_module, path, code_lines=[]),
                             ) for context in context_set
                         )
+                return None
+
+            context_set = callback(evaluator, import_names, parent_module_context, sys_path)
+            import_name = import_names[0]
+            if len(import_names) == 1 and import_name != 'typing':
+                map_ = self._cache_stub_file_map(evaluator.grammar.version_info)
+                result = _find_and_load_stub_module(map_)
+                if result is not None:
+                    return result
+            elif isinstance(parent_module_context, ModuleStubProxy):
+                map_ = _merge_create_stub_map(parent_module_context.stub_py__path__())
+                result = _find_and_load_stub_module(map_)
+                if result is not None:
+                    return result
             return context_set
         return wrapper
 
 
 class StubProxy(object):
-    def __init__(self, parent_context, context, stub_context):
-        self.parent_context = parent_context
+    def __init__(self, context, stub_context):
         self._context = context
         self._stub_context = stub_context
 
@@ -127,7 +144,7 @@ class StubProxy(object):
             return NO_CONTEXTS
 
         return ContextSet.from_iterable(
-            StubProxy(c.parent_context, c, typeshed_results[0]) for c in context_results
+            StubProxy(c, typeshed_results[0]) for c in context_results
         )
 
     @property
@@ -142,3 +159,12 @@ class StubProxy(object):
 
     def __repr__(self):
         return '<%s: %s %s>' % (type(self).__name__, self._context, self._stub_context)
+
+
+class ModuleStubProxy(StubProxy):
+    def __init__(self, parent_module_context, *args, **kwargs):
+        super(ModuleStubProxy, self).__init__(*args, **kwargs)
+        self._parent_module_context = parent_module_context
+
+    def stub_py__path__(self):
+        return self._stub_context.py__path__()