Add support for ZIP and EGG packages in imports

jedi/_compatibility.py

@@ -6,6 +6,7 @@ import sys
 import imp
 import os
 import re
+import pkgutil
 try:
     import importlib
 except ImportError:
@@ -18,6 +19,18 @@ is_py35 = is_py3 and sys.version_info.minor >= 5
 is_py26 = not is_py3 and sys.version_info[1] < 7
 
 
+class DummyFile(object):
+    def __init__(self, loader, string):
+        self.loader = loader
+        self.string = string
+
+    def read(self):
+        return self.loader.get_source(self.string)
+
+    def close(self):
+        del self.loader
+
+
 def find_module_py33(string, path=None):
     loader = importlib.machinery.PathFinder.find_module(string, path)
 
@@ -35,30 +48,73 @@ def find_module_py33(string, path=None):
     try:
         is_package = loader.is_package(string)
         if is_package:
-            module_path = os.path.dirname(loader.path)
-            module_file = None
+            if hasattr(loader, 'path'):
+                module_path = os.path.dirname(loader.path)
+            else:
+                # At least zipimporter does not have path attribute
+                module_path = os.path.dirname(loader.get_filename(string))
+            if hasattr(loader, 'archive'):
+                module_file = DummyFile(loader, string)
+            else:
+                module_file = None
         else:
             module_path = loader.get_filename(string)
-            module_file = open(module_path, 'rb')
+            module_file = DummyFile(loader, string)
     except AttributeError:
         # ExtensionLoader has not attribute get_filename, instead it has a
         # path attribute that we can use to retrieve the module path
         try:
             module_path = loader.path
-            module_file = open(loader.path, 'rb')
+            module_file = DummyFile(loader, string)
         except AttributeError:
             module_path = string
             module_file = None
         finally:
             is_package = False
 
+    if hasattr(loader, 'archive'):
+        module_path = loader.archive
+
     return module_file, module_path, is_package
 
 
 def find_module_pre_py33(string, path=None):
-    module_file, module_path, description = imp.find_module(string, path)
-    module_type = description[2]
-    return module_file, module_path, module_type is imp.PKG_DIRECTORY
+    try:
+        module_file, module_path, description = imp.find_module(string, path)
+        module_type = description[2]
+        return module_file, module_path, module_type is imp.PKG_DIRECTORY
+    except ImportError:
+        pass
+
+    if path is None:
+        path = sys.path
+    for item in path:
+        loader = pkgutil.get_importer(item)
+        if loader:
+            try:
+                loader = loader.find_module(string)
+                if loader:
+                    is_package = loader.is_package(string)
+                    is_archive = hasattr(loader, 'archive')
+                    try:
+                        module_path = loader.get_filename(string)
+                    except AttributeError:
+                        # fallback for py26
+                        try:
+                            module_path = loader._get_filename(string)
+                        except AttributeError:
+                            continue
+                    if is_package:
+                        module_path = os.path.dirname(module_path)
+                    if is_archive:
+                        module_path = loader.archive
+                    file = None
+                    if not is_package or is_archive:
+                        file = DummyFile(loader, string)
+                    return (file, module_path, is_package)
+            except ImportError:
+                pass
+    raise ImportError("No module named {0}".format(string))
 
 
 find_module = find_module_py33 if is_py33 else find_module_pre_py33

jedi/evaluate/imports.py

@@ -337,12 +337,15 @@ class Importer(object):
         if is_pkg:
             # In this case, we don't have a file yet. Search for the
             # __init__ file.
-            module_path = get_init_path(module_path)
+            if module_path.endswith(('.zip', '.egg')):
+                source = module_file.loader.get_source(module_name)
+            else:
+                module_path = get_init_path(module_path)
         elif module_file:
             source = module_file.read()
             module_file.close()
 
-        if module_file is None and not module_path.endswith('.py'):
+        if module_file is None and not module_path.endswith(('.py', '.zip', '.egg')):
             module = compiled.load_module(self._evaluator, module_path)
         else:
             module = _load_module(self._evaluator, module_path, source, sys_path)
@@ -444,7 +447,7 @@ class Importer(object):
 def _load_module(evaluator, path=None, source=None, sys_path=None):
     def load(source):
         dotted_path = path and compiled.dotted_from_fs_path(path, sys_path)
-        if path is not None and path.endswith('.py') \
+        if path is not None and path.endswith(('.py', '.zip', '.egg')) \
                 and dotted_path not in settings.auto_import_modules:
             if source is None:
                 with open(path, 'rb') as f:

test/test_evaluate/test_imports.py

@@ -4,7 +4,7 @@ import sys
 import pytest
 
 import jedi
-from jedi._compatibility import find_module_py33
+from jedi._compatibility import find_module_py33, find_module
 from ..helpers import cwd_at
 
 
@@ -14,6 +14,44 @@ def test_find_module_py33():
     assert find_module_py33('_io') == (None, '_io', False)
 
 
+def test_find_module_package():
+    file, path, is_package = find_module('json')
+    assert file is None
+    assert path.endswith('json')
+    assert is_package is True
+
+
+def test_find_module_not_package():
+    file, path, is_package = find_module('io')
+    assert file is not None
+    assert path.endswith('io.py')
+    assert is_package is False
+
+
+def test_find_module_package_zipped():
+    if 'zipped_imports/pkg.zip' not in sys.path:
+      sys.path.append(os.path.join(os.path.dirname(__file__),
+                      'zipped_imports/pkg.zip'))
+    file, path, is_package = find_module('pkg')
+    assert file is not None
+    assert path.endswith('pkg.zip')
+    assert is_package is True
+    assert len(jedi.Script('import pkg; pkg.mod', 1, 19).completions()) == 1
+
+
+@pytest.mark.skipif('sys.version_info < (2,7)')
+def test_find_module_not_package_zipped():
+    if 'zipped_imports/not_pkg.zip' not in sys.path:
+      sys.path.append(os.path.join(os.path.dirname(__file__),
+                      'zipped_imports/not_pkg.zip'))
+    file, path, is_package = find_module('not_pkg')
+    assert file is not None
+    assert path.endswith('not_pkg.zip')
+    assert is_package is False
+    assert len(
+      jedi.Script('import not_pkg; not_pkg.val', 1, 27).completions()) == 1
+
+
 @cwd_at('test/test_evaluate/not_in_sys_path/pkg')
 def test_import_not_in_sys_path():
     """