Additional helper methods, to find all the statement elements that are needed.

jedi/api/__init__.py

@@ -689,6 +689,11 @@ def defined_names(source, path=None, encoding='utf-8'):
 
 
 def _names(source=None, path=None, encoding='utf-8'):
+    """
+    Returns a list of `Definition` objects, containing name parts.
+    This means you can call ``Definition.goto_assignments()`` and get the
+    reference of a name.
+    """
     # Set line/column to a random position, because they don't matter.
     script = Script(source, line=1, column=0, path=path, encoding=encoding)
     defs = [classes.Definition(script._evaluator, name_part)

jedi/api/classes.py

@@ -11,6 +11,7 @@ from jedi import settings
 from jedi import common
 from jedi.parser import representation as pr
 from jedi.cache import underscore_memoization
+from jedi.evaluate.helpers import statement_elements_in_statement
 from jedi.evaluate.cache import memoize_default, CachedMetaClass
 from jedi.evaluate import representation as er
 from jedi.evaluate import iterable
@@ -310,6 +311,16 @@ class BaseDefinition(object):
         return '.'.join(path if path[0] else path[1:])
 
     def goto_assignments(self):
+        def call_path_for_name_part(stmt_or_imp, name_part):
+            if isinstance(stmt_or_imp, pr.Import):
+                return [name_part]
+            else:
+                for stmt_el in statement_elements_in_statement(stmt_or_imp):
+                    call_path = list(stmt_el.generate_call_path())
+                    for i, element in enumerate(call_path):
+                        if element is name_part:
+                            return call_path[:i+1]
+
         if not isinstance(self._definition, pr.NamePart):
             raise TypeError('Definition is not a NamePart.')
 
@@ -317,9 +328,9 @@ class BaseDefinition(object):
             # Functions, classes and modules are already fixed definitions, we
             # cannot follow them anymore.
             return [self]
-        stmt_or_imp = self._definition.parent.parent
-
-        names, _ = self._evaluator.goto(stmt_or_imp, [self._definition])
+        stmt_or_imp = self._definition.get_parent_until((pr.Statement, pr.Import))
+        call_path = call_path_for_name_part(stmt_or_imp, self._definition)
+        names, _ = self._evaluator.goto(stmt_or_imp, call_path)
         return [Definition(self._evaluator, n) for n in names]
 
     @memoize_default()

jedi/evaluate/helpers.py

@@ -107,6 +107,8 @@ def call_signature_array_for_pos(stmt, pos):
         tup = None, 0, None
         while call.next is not None and tup[0] is None:
             method = search_array if isinstance(call.next, pr.Array) else search_call
+            # TODO This is wrong, don't call search_call again, because it will
+            # automatically be called by call.next.
             tup = method(call.next, pos, origin_call or call)
             call = call.next
         return tup
@@ -224,6 +226,26 @@ def get_module_name_parts(module):
     return name_parts
 
 
+def statement_elements_in_statement(stmt):
+    """
+    Returns a list of statements. Statements can contain statements again in
+    Arrays.
+    """
+    def search_stmt_el(stmt_el, stmt_els):
+        stmt_els.append(stmt_el)
+        while stmt_el is not None:
+            if isinstance(stmt_el, pr.Array):
+                for stmt in stmt_el.values + stmt_el.keys:
+                    stmt_els.extend(statement_elements_in_statement(stmt))
+            stmt_el = stmt_el.next
+
+    stmt_els = []
+    for item in stmt.expression_list():
+        if isinstance(item, pr.StatementElement):
+            search_stmt_el(item, stmt_els)
+    return stmt_els
+
+
 class FakeSubModule():
     line_offset = 0