From e79ebe3ee7dd3fdd68c6fdbf0784f847bcd7b9f7 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Sat, 19 Nov 2016 03:05:10 +0100
Subject: [PATCH] Usage fixes for imports.

---
 jedi/api/__init__.py     | 14 ++++----------
 jedi/api/usages.py       | 40 ++++++++++++++++++++++------------------
 jedi/evaluate/imports.py |  7 ++++++-
 3 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/jedi/api/__init__.py b/jedi/api/__init__.py
index f7868d94..24087a42 100644
--- a/jedi/api/__init__.py
+++ b/jedi/api/__init__.py
@@ -266,28 +266,22 @@ class Script(object):
                 if name is None:
                     # Must be syntax
                     return []
-                self._evaluator
                 definition_names = [TreeNameDefinition(self._get_module(), name)]
 
             if not definition_names:
                 # Without a definition for a name we cannot find references.
                 return []
 
-            if not isinstance(user_stmt, tree.Import) and False:
-                # import case is looked at with add_import_name option
-                definition_names = usages.usages_add_import_modules(self._evaluator,
-                                                                    definition_names)
+            definition_names = usages.resolve_potential_imports(self._evaluator,
+                                                                definition_names)
 
             modules = set([d.get_root_context() for d in definition_names])
             modules.add(self._get_module())
-            names = usages.usages(self._evaluator, definition_names, modules)
-
-            for d in set(definition_names):
-                names.append(classes.Definition(self._evaluator, d))
+            definitions = usages.usages(self._evaluator, definition_names, modules)
         finally:
             settings.dynamic_flow_information = temp
 
-        return helpers.sorted_definitions(set(names))
+        return helpers.sorted_definitions(set(definitions))
 
     def call_signatures(self):
         """
diff --git a/jedi/api/usages.py b/jedi/api/usages.py
index b469b77a..a9785194 100644
--- a/jedi/api/usages.py
+++ b/jedi/api/usages.py
@@ -1,8 +1,8 @@
-from jedi._compatibility import unicode
 from jedi.api import classes
 from jedi.parser import tree
 from jedi.evaluate import imports
 from jedi.evaluate.filters import TreeNameDefinition
+from jedi.evaluate.representation import ModuleContext
 
 
 def usages(evaluator, definition_names, mods):
@@ -22,27 +22,31 @@ def usages(evaluator, definition_names, mods):
     search_name = list(definition_names)[0].string_name
     compare_definitions = compare_array(definition_names)
     mods = mods | set([d.get_root_context() for d in definition_names])
-    definitions = []
+    definition_names = set(definition_names)
     for m in imports.get_modules_containing_name(evaluator, mods, search_name):
-        for name_node in m.module_node.used_names.get(search_name, []):
-            context = evaluator.create_context(m, name_node)
-            result = evaluator.goto(context, name_node)
-            if [c for c in compare_array(result) if c in compare_definitions]:
-                name = TreeNameDefinition(context, name_node)
-                definitions.append(classes.Definition(evaluator, name))
-                # Previous definitions might be imports, so include them
-                # (because goto might return that import name).
-                compare_definitions += compare_array([name])
-    return definitions
+        if isinstance(m, ModuleContext):
+            for name_node in m.module_node.used_names.get(search_name, []):
+                context = evaluator.create_context(m, name_node)
+                result = evaluator.goto(context, name_node)
+                if [c for c in compare_array(result) if c in compare_definitions]:
+                    name = TreeNameDefinition(context, name_node)
+                    definition_names.add(name)
+                    # Previous definitions might be imports, so include them
+                    # (because goto might return that import name).
+                    compare_definitions += compare_array([name])
+        else:
+            definition_names.add(m.name)
+
+    return [classes.Definition(evaluator, n) for n in definition_names]
 
 
-def usages_add_import_modules(evaluator, definitions):
+def resolve_potential_imports(evaluator, definitions):
     """ Adds the modules of the imports """
     new = set()
     for d in definitions:
-        print(d)
-        imp_or_stmt = d.get_definition()
-        if isinstance(imp_or_stmt, tree.Import):
-            s = imports.ImportWrapper(context, d)
-            new |= set(s.follow(is_goto=True))
+        if isinstance(d, TreeNameDefinition):
+            imp_or_stmt = d.tree_name.get_definition()
+            if isinstance(imp_or_stmt, tree.Import):
+                s = imports.ImportWrapper(d.parent_context, d.tree_name)
+                new |= resolve_potential_imports(evaluator, set(s.follow(is_goto=True)))
     return set(definitions) | new
diff --git a/jedi/evaluate/imports.py b/jedi/evaluate/imports.py
index 5bc6c164..2c9bdc9a 100644
--- a/jedi/evaluate/imports.py
+++ b/jedi/evaluate/imports.py
@@ -501,7 +501,12 @@ def get_modules_containing_name(evaluator, modules, name):
     # skip non python modules
     used_mod_paths = set()
     for m in modules:
-        used_mod_paths.add(m.py__file__())
+        try:
+            path = m.py__file__()
+        except AttributeError:
+            pass
+        else:
+            used_mod_paths.add(path)
         yield m
 
     if not settings.dynamic_params_for_other_modules: