diff --git a/jedi/evaluate/context/module.py b/jedi/evaluate/context/module.py index 03fb0e55..5b4aab4e 100644 --- a/jedi/evaluate/context/module.py +++ b/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, diff --git a/jedi/evaluate/imports.py b/jedi/evaluate/imports.py index 34f8053e..5df13225 100644 --- a/jedi/evaluate/imports.py +++ b/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) diff --git a/jedi/plugins/stdlib.py b/jedi/plugins/stdlib.py index 4c1a8666..a47efddb 100644 --- a/jedi/plugins/stdlib.py +++ b/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 diff --git a/jedi/plugins/typeshed.py b/jedi/plugins/typeshed.py index 6091642a..57f7f19e 100644 --- a/jedi/plugins/typeshed.py +++ b/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__()