diff --git a/jedi/api/classes.py b/jedi/api/classes.py index 00dbddcc..b0f44210 100644 --- a/jedi/api/classes.py +++ b/jedi/api/classes.py @@ -50,6 +50,7 @@ class BaseDefinition(object): '_io': 'io', '_functools': 'functools', '_collections': 'collections', + '_socket': 'socket', '_sqlite3': 'sqlite3', '__builtin__': 'builtins', } diff --git a/jedi/evaluate/gradual/typeshed.py b/jedi/evaluate/gradual/typeshed.py index 22b855ca..1fc8d323 100644 --- a/jedi/evaluate/gradual/typeshed.py +++ b/jedi/evaluate/gradual/typeshed.py @@ -10,6 +10,11 @@ from jedi.evaluate.gradual.stub_context import TypingModuleWrapper, StubModuleCo _jedi_path = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) TYPESHED_PATH = os.path.join(_jedi_path, 'third_party', 'typeshed') +_IMPORT_MAP = dict( + _collections='collections', + _socket='socket', +) + def _merge_create_stub_map(directories): map_ = {} @@ -227,8 +232,7 @@ def _load_from_typeshed(evaluator, python_context_set, parent_module_context, im map_ = None if len(import_names) == 1: map_ = _cache_stub_file_map(evaluator.grammar.version_info) - if import_name == '_collections': - import_name = 'collections' + import_name = _IMPORT_MAP.get(import_name, import_name) elif isinstance(parent_module_context, StubModuleContext): if not parent_module_context.is_package: # Only if it's a package (= a folder) something can be diff --git a/test/test_evaluate/test_gradual/test_stubs.py b/test/test_evaluate/test_gradual/test_stubs.py index 35da0add..4e58c881 100644 --- a/test/test_evaluate/test_gradual/test_stubs.py +++ b/test/test_evaluate/test_gradual/test_stubs.py @@ -31,6 +31,9 @@ from test.helpers import root_dir ['from keyword import kwlist', 'typing.Sequence', True, True, {'full_name': 'keyword.kwlist'}], + ['from socket import AF_INET', 'socket.AddressFamily', True, True, + {'full_name': 'socket.AF_INET'}], + ['import with_stub', 'with_stub', True, True, {}], ['import with_stub', 'with_stub', True, True, {}], ['import with_stub_folder.python_only', 'with_stub_folder.python_only', False, True, {}],