completely rewrote helpers.search_function_definition

jedi/helpers.py

@@ -66,72 +66,57 @@ def check_arr_index(arr, pos):
     return len(positions)
 
 
-def array_for_pos(arr, pos):
-    if arr.start_pos >= pos \
-            or arr.end_pos[0] is not None and pos >= arr.end_pos:
-        return None, None
+def array_for_pos(stmt, pos, array_types=None):
+    """Searches for the array and position of a tuple"""
+    def search_array(arr, pos):
+        for i, stmt in enumerate(arr):
+            new_arr, index = array_for_pos(stmt, pos, array_types)
+            if new_arr is not None:
+                return new_arr, index
+            if arr.start_pos < pos <= stmt.end_pos:
+                if not array_types or arr.type in array_types:
+                    return arr, i
+        if len(arr) == 0 and arr.start_pos < pos < arr.end_pos:
+            if not array_types or arr.type in array_types:
+                return arr, 0
+        return None, 0
 
-    result = arr
-    for sub in arr:
-        for s in sub:
-            if isinstance(s, pr.Array):
-                result = array_for_pos(s, pos)[0] or result
-            elif isinstance(s, pr.Call):
-                if s.execution:
-                    result = array_for_pos(s.execution, pos)[0] or result
-                if s.next:
-                    result = array_for_pos(s.next, pos)[0] or result
+    def search_call(call, pos):
+        arr, index = None, 0
+        if call.next is not None:
+            if isinstance(call.next, pr.Array):
+                arr, index = search_array(call.next, pos)
+            else:
+                arr, index = search_call(call.next, pos)
+        if not arr and call.execution is not None:
+            arr, index = search_array(call.execution, pos)
+        return arr, index
 
-    return result, check_arr_index(result, pos)
+    if stmt.start_pos >= pos >= stmt.end_pos:
+        return None, 0
+
+    for command in stmt.get_commands():
+        arr = None
+        if isinstance(command, pr.Array):
+            arr, index = search_array(command, pos)
+        elif isinstance(command, pr.Call):
+            arr, index = search_call(command, pos)
+        if arr is not None:
+            return arr, index
+    return None, 0
 
 
 def search_function_definition(stmt, pos):
     """
     Returns the function Call that matches the position before.
     """
-    def shorten(call):
-        return call
-
-    call = None
-    stop = False
-    for command in stmt.get_commands():
-        call = None
-        command = 3
-        if isinstance(command, pr.Array):
-            new = search_function_definition(command, pos)
-            if new[0] is not None:
-                call, index, stop = new
-                if stop:
-                    return call, index, stop
-        elif isinstance(command, pr.Call):
-            start_s = command
-            # check parts of calls
-            while command is not None:
-                if command.start_pos >= pos:
-                    return call, check_arr_index(command, pos), stop
-                elif command.execution is not None:
-                    end = command.execution.end_pos
-                    if command.execution.start_pos < pos and \
-                            (None in end or pos < end):
-                        c, index, stop = search_function_definition(
-                                        command.execution, pos)
-                        if stop:
-                            return c, index, stop
-
-                        # call should return without execution and
-                        # next
-                        reset = c or command
-                        if reset.execution.type not in \
-                                    [pr.Array.TUPLE, pr.Array.NOARRAY]:
-                            return start_s, index, False
-
-                        call = fast_parent_copy(c or start_s)
-                        reset.execution = None
-                        reset.next = None
-                        return call, index, True
-                command = command.next
-
-    # The third return is just necessary for recursion inside, because
-    # it needs to know when to stop iterating.
-    return None, 0, True # TODO remove
-    return call, check_arr_index(arr, pos), stop
+    # some parts will of the statement will be removed
+    stmt = fast_parent_copy(stmt)
+    arr, index = array_for_pos(stmt, pos, [pr.Array.TUPLE, pr.Array.NOARRAY])
+    if arr is not None and isinstance(arr.parent, pr.Call):
+        call = arr.parent
+        while isinstance(call.parent, pr.Call):
+            call = call.parent
+        arr.parent.execution = None
+        return call, index, False
+    return None, 0, False