Try to change the module cache

jedi/evaluate/__init__.py

@@ -62,6 +62,7 @@ I need to mention now that lazy evaluation is really good because it
 only *evaluates* what needs to be *evaluated*. All the statements and modules
 that are not used are just being ignored.
 """
+from functools import partial
 
 from parso.python import tree
 import parso
@@ -125,11 +126,24 @@ class Evaluator(object):
         from jedi.plugins import plugin_manager
         plugin_callbacks = plugin_manager.get_callbacks(self)
         self.execute = plugin_callbacks.decorate('execute', callback=_execute)
-        self.import_module = plugin_callbacks.decorate(
-            'import_module',
-            callback=imports.import_module
+        self._import_module = partial(
+            plugin_callbacks.decorate(
+                'import_module',
+                callback=imports.import_module
+            ),
+            self,
         )
 
+    def import_module(self, import_names, parent_module_context, sys_path):
+        try:
+            return ContextSet(self.module_cache.get(import_names))
+        except KeyError:
+            pass
+
+        context_set = self._import_module(import_names, parent_module_context, sys_path)
+        self.module_cache.add(context_set, import_names)
+        return context_set
+
     @property
     @evaluator_function_cache()
     def builtins_module(self):

jedi/evaluate/compiled/mixed.py

@@ -203,6 +203,7 @@ def _create(evaluator, access_handle, parent_context, *args):
     if parent_context.tree_node.get_root_node() == module_node:
         module_context = parent_context.get_root_context()
     else:
+        raise NotImplementedError('module misses string_names arg')
         module_context = ModuleContext(
             evaluator, module_node,
             path=path,

jedi/evaluate/context/module.py

@@ -42,13 +42,14 @@ class ModuleContext(TreeContext):
     api_type = u'module'
     parent_context = None
 
-    def __init__(self, evaluator, module_node, path, code_lines):
+    def __init__(self, evaluator, module_node, path, string_names, code_lines):
         super(ModuleContext, self).__init__(
             evaluator,
             parent_context=None,
             tree_node=module_node
         )
         self._path = path
+        self._string_names = string_names
         self.code_lines = code_lines
 
     def get_filters(self, search_global, until_position=None, origin_scope=None):
@@ -120,11 +121,7 @@ class ModuleContext(TreeContext):
         return None
 
     def py__name__(self):
-        for name, module in self.evaluator.module_cache.iterate_modules_with_names():
-            if module == self and name != '':
-                return name
-
-        return '__main__'
+        return '.'.join(self._string_names)
 
     def py__file__(self):
         """

jedi/evaluate/imports.py

@@ -37,16 +37,14 @@ class ModuleCache(object):
         self._path_cache = {}
         self._name_cache = {}
 
-    def add(self, module, name):
-        path = module.py__file__()
-        self._path_cache[path] = module
-        self._name_cache[name] = module
+    def add(self, string_names, context_set):
+        #path = module.py__file__()
+        #self._path_cache[path] = context_set
+        if string_names is not None:
+            self._name_cache[string_names] = context_set
 
-    def iterate_modules_with_names(self):
-        return self._name_cache.items()
-
-    def get(self, name):
-        return self._name_cache[name]
+    def get(self, string_names):
+        return self._name_cache[string_names]
 
     def get_from_path(self, path):
         return self._path_cache[path]
@@ -291,7 +289,6 @@ class Importer(object):
             try:
                 context_set = ContextSet.from_sets([
                     self._evaluator.import_module(
-                        self._evaluator,
                         import_names[:i+1],
                         parent_module_context,
                         self.sys_path_with_modifications(),
@@ -410,11 +407,6 @@ def import_module(evaluator, import_names, parent_module_context, sys_path):
         return ContextSet(module)
 
     module_name = '.'.join(import_names)
-    try:
-        return ContextSet(evaluator.module_cache.get(module_name))
-    except KeyError:
-        pass
-
     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.
@@ -467,10 +459,6 @@ def _load_module(evaluator, path=None, code=None, sys_path=None,
         dotted_name = None
     else:
         dotted_name = '.'.join(import_names)
-    try:
-        return evaluator.module_cache.get(dotted_name)
-    except KeyError:
-        pass
     try:
         return evaluator.module_cache.get_from_path(path)
     except KeyError:
@@ -497,6 +485,7 @@ def _load_module(evaluator, path=None, code=None, sys_path=None,
             module = ModuleContext(
                 evaluator, module_node,
                 path=path,
+                string_names=import_names,
                 code_lines=get_cached_code_lines(evaluator.grammar, path),
             )
         else:
@@ -506,23 +495,9 @@ def _load_module(evaluator, path=None, code=None, sys_path=None,
                 # The file might raise an ImportError e.g. and therefore not be
                 # importable.
                 raise JediImportError(import_names)
-
-    if module is not None and dotted_name is not None:
-        add_module_to_cache(evaluator, dotted_name, module, safe=safe_module_name)
-
     return module
 
 
-def add_module_to_cache(evaluator, module_name, module, safe=False):
-    if not safe and '.' not in module_name:
-        # We cannot add paths with dots, because that would collide with
-        # the sepatator dots for nested packages. Therefore we return
-        # `__main__` in ModuleWrapper.py__name__(), which is similar to
-        # Python behavior.
-        return
-    evaluator.module_cache.add(module, module_name)
-
-
 def get_modules_containing_name(evaluator, modules, name):
     """
     Search a name in the directories of modules.
@@ -553,6 +528,7 @@ def get_modules_containing_name(evaluator, modules, name):
                     sys_path=e_sys_path,
                     import_names=import_names,
                 )
+                evaluator.module_cache.add(import_names, ContextSet(module))
                 return module
 
     # skip non python modules

jedi/evaluate/sys_path.py

@@ -151,6 +151,7 @@ def _get_paths_from_buildout_script(evaluator, buildout_script_path):
         return
 
     from jedi.evaluate.context import ModuleContext
+    raise NotImplementedError("No module string_names, yet")
     module = ModuleContext(
         evaluator, module_node, buildout_script_path,
         code_lines=get_cached_code_lines(evaluator.grammar, buildout_script_path),