huge functions refactoring

functions.py

@@ -12,8 +12,7 @@ import keywords
 
 from _compatibility import next
 
-__all__ = ['complete', 'goto', 'get_definition', 'related_names',
-           'NotFoundError', 'set_debug_function', 'get_in_function_call']
+__all__ = ['Script', 'NotFoundError', 'set_debug_function']
 
 
 class NotFoundError(Exception):
@@ -133,249 +132,292 @@ class CallDef(object):
                                     self.index)
 
 
-def _get_completion_parts(path):
-    """
-    Returns the parts for the completion
-    :return: tuple - (path, dot, like)
-    """
-    match = re.match(r'^(.*?)(\.|)(\w?[\w\d]*)$', path, flags=re.S)
-    return match.groups()
+class Script(object):
+    """ TODO doc """
+    def __init__(self, source, line, column, source_path):
+        self.pos = line, column
+        self.module = modules.ModuleWithCursor(source_path, source=source,
+                                                            position=self.pos)
+        self.parser = self.module.parser
+        self.source_path = source_path
 
+    def complete(self):
+        """
+        An auto completer for python files.
 
-def complete(source, line, column, source_path):
-    """
-    An auto completer for python files.
+        :param source: The source code of the current file
+        :type source: string
+        :param line: The line to complete in.
+        :type line: int
+        :param col: The column to complete in.
+        :type col: int
+        :param source_path: The path in the os, the current module is in.
+        :type source_path: str
 
-    :param source: The source code of the current file
-    :type source: string
-    :param line: The line to complete in.
-    :type line: int
-    :param col: The column to complete in.
-    :type col: int
-    :param source_path: The path in the os, the current module is in.
-    :type source_path: str
+        :return: list of Completion objects.
+        :rtype: list
+        """
+        path = self.module.get_path_until_cursor()
+        path, dot, like = self._get_completion_parts(path)
 
-    :return: list of Completion objects.
-    :rtype: list
-    """
-    pos = (line, column)
-    f = modules.ModuleWithCursor(source_path, source=source, position=pos)
-    path = f.get_path_until_cursor()
-    path, dot, like = _get_completion_parts(path)
-
-    try:
-        scopes = _prepare_goto(pos, source_path, f, path, True)
-    except NotFoundError:
-        scope_generator = evaluate.get_names_for_scope(f.parser.user_scope,
-                                                                        pos)
-        completions = []
-        for scope, name_list in scope_generator:
-            for c in name_list:
-                completions.append((c, scope))
-    else:
-        completions = []
-        debug.dbg('possible scopes', scopes)
-        for s in scopes:
-            # TODO is this really the right way? just ignore the functions? \
-            # do the magic functions first? and then recheck here?
-            if not isinstance(s, evaluate.Function):
-                if isinstance(s, imports.ImportPath):
-                    names = s.get_defined_names(on_import_stmt=True)
-                else:
-                    names = s.get_defined_names()
-                for c in names:
-                    completions.append((c, s))
-
-    completions = [(c, s) for c, s in completions
-                        if settings.case_insensitive_completion
-                            and c.names[-1].lower().startswith(like.lower())
-                            or c.names[-1].startswith(like)]
-
-    needs_dot = not dot and path
-    c = [Completion(c, needs_dot, len(like), s) for c, s in set(completions)]
-
-    call_def = _get_in_function_call(f, pos)
-
-    _clear_caches()
-    return c, call_def
-
-
-def _prepare_goto(position, source_path, module, goto_path,
-                                                        is_like_search=False):
-    scope = module.parser.user_scope
-    debug.dbg('start: %s in %s' % (goto_path, scope))
-
-    user_stmt = module.parser.user_stmt
-    if not user_stmt and len(goto_path.split('\n')) > 1:
-        # If the user_stmt is not defined and the goto_path is multi line,
-        # something's strange. Most probably the backwards tokenizer matched to
-        # much.
-        return []
-
-    if isinstance(user_stmt, parsing.Import):
-        import_names = user_stmt.get_all_import_names()
-        count = 0
-        kill_count = -1
-        for i in import_names:
-            for name_part in i.names:
-                count += 1
-                if position <= name_part.end_pos:
-                    kill_count += 1
-        scopes = [imports.ImportPath(user_stmt, is_like_search,
-                        kill_count=kill_count, direct_resolve=True)]
-    else:
-        # just parse one statement, take it and evaluate it
-        r = parsing.PyFuzzyParser(goto_path, source_path, no_docstr=True)
         try:
-            stmt = r.module.statements[0]
-        except IndexError:
-            raise NotFoundError()
+            scopes = self._prepare_goto(path, True)
+        except NotFoundError:
+            scope_generator = evaluate.get_names_for_scope(self.parser.user_scope,
+                                                                            self.pos)
+            completions = []
+            for scope, name_list in scope_generator:
+                for c in name_list:
+                    completions.append((c, scope))
+        else:
+            completions = []
+            debug.dbg('possible scopes', scopes)
+            for s in scopes:
+                # TODO is this really the right way? just ignore the functions? \
+                # do the magic functions first? and then recheck here?
+                if not isinstance(s, evaluate.Function):
+                    if isinstance(s, imports.ImportPath):
+                        names = s.get_defined_names(on_import_stmt=True)
+                    else:
+                        names = s.get_defined_names()
+                    for c in names:
+                        completions.append((c, s))
 
-        stmt.start_pos = position
-        stmt.parent = weakref.ref(scope)
-        scopes = evaluate.follow_statement(stmt)
-    return scopes
+        completions = [(c, s) for c, s in completions
+                            if settings.case_insensitive_completion
+                                and c.names[-1].lower().startswith(like.lower())
+                                or c.names[-1].startswith(like)]
 
+        needs_dot = not dot and path
+        c = [Completion(c, needs_dot, len(like), s) for c, s in set(completions)]
 
-def get_definition(source, line, column, source_path):
-    """
-    Returns the definitions of a the path under the cursor.
-    This is not a goto function! This follows complicated paths and returns the
-    end, not the first definition.
+        _clear_caches()
+        return c
 
-    :param source: The source code of the current file
-    :type source: string
-    :param line: The line to complete in.
-    :type line: int
-    :param col: The column to complete in.
-    :type col: int
-    :param source_path: The path in the os, the current module is in.
-    :type source_path: int
+    def _prepare_goto(self, goto_path, is_like_search=False):
+        scope = self.parser.user_scope
+        debug.dbg('start: %s in %s' % (goto_path, scope))
 
-    :return: list of Definition objects, which are basically scopes.
-    :rtype: list
-    """
-    def resolve_import_paths(scopes):
-        for s in scopes.copy():
-            if isinstance(s, imports.ImportPath):
-                scopes.remove(s)
-                scopes.update(resolve_import_paths(set(s.follow())))
+        user_stmt = self.parser.user_stmt
+        if not user_stmt and len(goto_path.split('\n')) > 1:
+            # If the user_stmt is not defined and the goto_path is multi line,
+            # something's strange. Most probably the backwards tokenizer matched to
+            # much.
+            return []
+
+        if isinstance(user_stmt, parsing.Import):
+            import_names = user_stmt.get_all_import_names()
+            count = 0
+            kill_count = -1
+            for i in import_names:
+                for name_part in i.names:
+                    count += 1
+                    if self.pos <= name_part.end_pos:
+                        kill_count += 1
+            scopes = [imports.ImportPath(user_stmt, is_like_search,
+                            kill_count=kill_count, direct_resolve=True)]
+        else:
+            # just parse one statement, take it and evaluate it
+            r = parsing.PyFuzzyParser(goto_path, self.source_path, no_docstr=True)
+            try:
+                stmt = r.module.statements[0]
+            except IndexError:
+                raise NotFoundError()
+
+            stmt.start_pos = self.pos
+            stmt.parent = weakref.ref(scope)
+            scopes = evaluate.follow_statement(stmt)
         return scopes
 
-    pos = (line, column)
-    f = modules.ModuleWithCursor(source_path, source=source, position=pos)
-    goto_path = f.get_path_under_cursor()
 
-    context = f.get_context()
-    if next(context) in ('class', 'def'):
-        scopes = set([f.parser.user_scope])
-    elif not goto_path:
-        op = f.get_operator_under_cursor()
-        scopes = set([keywords.get_operator(op, pos)] if op else [])
-    else:
-        scopes = set(_prepare_goto(pos, source_path, f, goto_path))
+    def get_definition(self):
+        """
+        Returns the definitions of a the path under the cursor.
+        This is not a goto function! This follows complicated paths and returns the
+        end, not the first definition.
 
-    scopes = resolve_import_paths(scopes)
+        :param source: The source code of the current file
+        :type source: string
+        :param line: The line to complete in.
+        :type line: int
+        :param col: The column to complete in.
+        :type col: int
+        :param source_path: The path in the os, the current module is in.
+        :type source_path: int
 
-    # add keywords
-    scopes |= keywords.get_keywords(string=goto_path, pos=pos)
+        :return: list of Definition objects, which are basically scopes.
+        :rtype: list
+        """
+        def resolve_import_paths(scopes):
+            for s in scopes.copy():
+                if isinstance(s, imports.ImportPath):
+                    scopes.remove(s)
+                    scopes.update(resolve_import_paths(set(s.follow())))
+            return scopes
 
-    d = set([Definition(s) for s in scopes])
-    _clear_caches()
-    return sorted(d, key=lambda x: (x.module_path, x.start_pos))
+        goto_path = self.module.get_path_under_cursor()
 
-
-def goto(source, line, column, source_path):
-    pos = (line, column)
-    f = modules.ModuleWithCursor(source_path, source=source, position=pos)
-
-    goto_path = f.get_path_under_cursor()
-    goto_path, dot, search_name = _get_completion_parts(goto_path)
-
-    # define goto path the right way
-    if not dot:
-        goto_path = search_name
-        search_name_new = None
-    else:
-        search_name_new = search_name
-
-    context = f.get_context()
-    if next(context) in ('class', 'def'):
-        definitions = set([f.parser.user_scope])
-    else:
-        scopes = _prepare_goto(pos, source_path, f, goto_path)
-        definitions = evaluate.goto(scopes, search_name_new)
-
-    d = [Definition(d) for d in set(definitions)]
-    _clear_caches()
-    return sorted(d, key=lambda x: (x.module_path, x.start_pos))
-
-
-def related_names(source, line, column, source_path):
-    """
-    Returns `dynamic.RelatedName` objects, which contain all names, that are
-    defined by the same variable, function, class or import.
-    This function can be used either to show all the usages of a variable or
-    for renaming purposes.
-    """
-    pos = (line, column)
-    f = modules.ModuleWithCursor(source_path, source=source, position=pos)
-
-    goto_path = f.get_path_under_cursor()
-    goto_path, dot, search_name = _get_completion_parts(goto_path)
-
-    # define goto path the right way
-    if not dot:
-        goto_path = search_name
-        search_name_new = None
-    else:
-        search_name_new = search_name
-
-    context = f.get_context()
-    if next(context) in ('class', 'def'):
-        if isinstance(f.parser.user_scope, parsing.Function):
-            e = evaluate.Function(f.parser.user_scope)
+        context = self.module.get_context()
+        if next(context) in ('class', 'def'):
+            scopes = set([self.module.parser.user_scope])
+        elif not goto_path:
+            op = self.module.get_operator_under_cursor()
+            scopes = set([keywords.get_operator(op, self.pos)] if op else [])
         else:
-            e = evaluate.Class(f.parser.user_scope)
-        definitions = [e]
-    elif isinstance(f.parser.user_stmt, (parsing.Param, parsing.Import)):
-        definitions = [f.parser.user_stmt]
-    else:
-        scopes = _prepare_goto(pos, source_path, f, goto_path)
-        definitions = evaluate.goto(scopes, search_name_new)
+            scopes = set(self._prepare_goto(goto_path))
 
-    module = set([d.get_parent_until() for d in definitions])
-    module.add(f.parser.module)
-    names = dynamic.related_names(definitions, search_name, module)
+        scopes = resolve_import_paths(scopes)
 
-    for d in definitions:
-        if isinstance(d, parsing.Statement):
-            def add_array(arr):
-                calls = dynamic._scan_array(arr, search_name)
-                for call in calls:
-                    for n in call.name.names:
-                        if n == search_name:
-                            names.append(dynamic.RelatedName(n, d))
-            for op, arr in d.assignment_details:
-                add_array(arr)
-            if not d.assignment_details:
-                add_array(d.get_assignment_calls())
-        elif isinstance(d, parsing.Import):
-            is_user = d == f.parser.user_stmt
-            check_names = [d.namespace, d.alias, d.from_ns] if is_user \
-                                                    else d.get_defined_names()
-            for name in check_names:
-                if name:
-                    for n in name.names:
-                        if n.start_pos <= pos <= n.end_pos or not is_user:
-                            names.append(dynamic.RelatedName(n, d))
-        elif isinstance(d, parsing.Name):
-            names.append(dynamic.RelatedName(d.names[0], d))
+        # add keywords
+        scopes |= keywords.get_keywords(string=goto_path, pos=self.pos)
+
+        d = set([Definition(s) for s in scopes])
+        _clear_caches()
+        return sorted(d, key=lambda x: (x.module_path, x.start_pos))
+
+
+    def goto(self):
+        goto_path = self.module.get_path_under_cursor()
+        goto_path, dot, search_name = self._get_completion_parts(goto_path)
+
+        # define goto path the right way
+        if not dot:
+            goto_path = search_name
+            search_name_new = None
         else:
-            names.append(dynamic.RelatedName(d.name.names[0], d))
+            search_name_new = search_name
 
-    _clear_caches()
-    return sorted(names, key=lambda x: (x.module_path, x.start_pos))
+        context = self.module.get_context()
+        if next(context) in ('class', 'def'):
+            definitions = set([self.module.parser.user_scope])
+        else:
+            scopes = self._prepare_goto(goto_path)
+            definitions = evaluate.goto(scopes, search_name_new)
+
+        d = [Definition(d) for d in set(definitions)]
+        _clear_caches()
+        return sorted(d, key=lambda x: (x.module_path, x.start_pos))
+
+
+    def related_names(self):
+        """
+        Returns `dynamic.RelatedName` objects, which contain all names, that are
+        defined by the same variable, function, class or import.
+        This function can be used either to show all the usages of a variable or
+        for renaming purposes.
+        """
+        goto_path = self.module.get_path_under_cursor()
+        goto_path, dot, search_name = self._get_completion_parts(goto_path)
+
+        # define goto path the right way
+        if not dot:
+            goto_path = search_name
+            search_name_new = None
+        else:
+            search_name_new = search_name
+
+        context = self.module.get_context()
+        if next(context) in ('class', 'def'):
+            if isinstance(self.module.parser.user_scope, parsing.Function):
+                e = evaluate.Function(self.module.parser.user_scope)
+            else:
+                e = evaluate.Class(self.module.parser.user_scope)
+            definitions = [e]
+        elif isinstance(self.module.parser.user_stmt, (parsing.Param, parsing.Import)):
+            definitions = [self.module.parser.user_stmt]
+        else:
+            scopes = self._prepare_goto(goto_path)
+            definitions = evaluate.goto(scopes, search_name_new)
+
+        module = set([d.get_parent_until() for d in definitions])
+        module.add(self.module.parser.module)
+        names = dynamic.related_names(definitions, search_name, module)
+
+        for d in definitions:
+            if isinstance(d, parsing.Statement):
+                def add_array(arr):
+                    calls = dynamic._scan_array(arr, search_name)
+                    for call in calls:
+                        for n in call.name.names:
+                            if n == search_name:
+                                names.append(dynamic.RelatedName(n, d))
+                for op, arr in d.assignment_details:
+                    add_array(arr)
+                if not d.assignment_details:
+                    add_array(d.get_assignment_calls())
+            elif isinstance(d, parsing.Import):
+                is_user = d == self.module.parser.user_stmt
+                check_names = [d.namespace, d.alias, d.from_ns] if is_user \
+                                                        else d.get_defined_names()
+                for name in check_names:
+                    if name:
+                        for n in name.names:
+                            if n.start_pos <= self.pos <= n.end_pos or not is_user:
+                                names.append(dynamic.RelatedName(n, d))
+            elif isinstance(d, parsing.Name):
+                names.append(dynamic.RelatedName(d.names[0], d))
+            else:
+                names.append(dynamic.RelatedName(d.name.names[0], d))
+
+        _clear_caches()
+        return sorted(names, key=lambda x: (x.module_path, x.start_pos))
+
+    def get_in_function_call(self):
+        def scan_array_for_pos(arr, pos):
+            """ Returns the function Call that match search_name in an Array. """
+            index = None
+            call = None
+            for index, sub in enumerate(arr):
+                call = None
+                for s in sub:
+                    if isinstance(s, parsing.Array):
+                        new = scan_array_for_pos(s, pos)
+                        if new[0] is not None:
+                            call, index = new
+                    elif isinstance(s, parsing.Call):
+                        while s is not None:
+                            if s.start_pos >= pos:
+                                return call, index
+                            if s.execution is not None:
+                                if s.execution.start_pos <= pos:
+                                    call = s
+                                    c, index = scan_array_for_pos(s.execution, pos)
+                                    if c is not None:
+                                        call = c
+                                else:
+                                    return call, index
+                            s = s.next
+            return call, index
+
+        user_stmt = self.parser.user_stmt
+        if user_stmt is None or not isinstance(user_stmt, parsing.Statement):
+            return None
+        ass = user_stmt.get_assignment_calls()
+
+        call, index = scan_array_for_pos(ass, self.pos)
+        if call is None:
+            return None
+
+        call.execution, temp = None, call.execution
+        origins = evaluate.follow_call(call)
+        call.execution = temp
+
+        if len(origins) == 0:
+            return None
+        executable = origins[0]  # just take entry zero, because we need just one.
+        return CallDef(executable, index)
+
+    def _get_completion_parts(self, path):
+        """
+        Returns the parts for the completion
+        :return: tuple - (path, dot, like)
+        """
+        match = re.match(r'^(.*?)(\.|)(\w?[\w\d]*)$', path, flags=re.S)
+        return match.groups()
+
+
+def _clear_caches():
+    evaluate.clear_caches()
 
 
 def set_debug_function(func_cb):
@@ -384,58 +426,3 @@ def set_debug_function(func_cb):
     :param func_cb: The callback function for debug messages, with n params.
     """
     debug.debug_function = func_cb
-
-
-def _clear_caches():
-    evaluate.clear_caches()
-
-def get_in_function_call(source, line, column, source_path):
-    pos = (line, column)
-    f = modules.ModuleWithCursor(source_path, source=source, position=pos)
-
-    return _get_in_function_call(f, pos)
-
-def _get_in_function_call(module, pos):
-    def scan_array_for_pos(arr, pos):
-        """ Returns the function Call that match search_name in an Array. """
-        index = None
-        call = None
-        for index, sub in enumerate(arr):
-            call = None
-            for s in sub:
-                if isinstance(s, parsing.Array):
-                    new = scan_array_for_pos(s, pos)
-                    if new[0] is not None:
-                        call, index = new
-                elif isinstance(s, parsing.Call):
-                    while s is not None:
-                        if s.start_pos >= pos:
-                            return call, index
-                        if s.execution is not None:
-                            if s.execution.start_pos <= pos:
-                                call = s
-                                c, index = scan_array_for_pos(s.execution, pos)
-                                if c is not None:
-                                    call = c
-                            else:
-                                return call, index
-                        s = s.next
-        return call, index
-
-    user_stmt = module.parser.user_stmt
-    if user_stmt is None or not isinstance(user_stmt, parsing.Statement):
-        return None
-    ass = user_stmt.get_assignment_calls()
-
-    call, index = scan_array_for_pos(ass, pos)
-    if call is None:
-        return None
-
-    call.execution, temp = None, call.execution
-    origins = evaluate.follow_call(call)
-    call.execution = temp
-
-    if len(origins) == 0:
-        return None
-    executable = origins[0]  # just take entry zero, because we need just one.
-    return CallDef(executable, index)

plugin/jedi.vim

@@ -16,10 +16,10 @@ if exists("g:loaded_jedi") || &cp
 endif
 let g:loaded_jedi = 1
 
+
 " ------------------------------------------------------------------------
 " completion
 " ------------------------------------------------------------------------
-
 function! jedi#complete(findstart, base)
 python << PYTHONEOF
 if 1:
@@ -32,7 +32,6 @@ if 1:
             count += 1
         vim.command('return %i' % (column - count))
     else:
-        buf_path = vim.current.buffer.name
         base = vim.eval('a:base')
         source = ''
         for i, line in enumerate(vim.current.buffer):
@@ -45,7 +44,10 @@ if 1:
         # here again, the hacks, because jedi has a different interface than vim
         column += len(base)
         try:
-            completions, call_def = functions.complete(source, row, column, buf_path)
+            script = get_script(source=source, column=column)
+            completions = script.complete()
+            call_def = script.get_in_function_call()
+
             out = []
             for c in completions:
                 d = dict(word=c.word[:len(base)] + c.complete,
@@ -72,21 +74,16 @@ if 1:
 PYTHONEOF
 endfunction
 
+
 " ------------------------------------------------------------------------
 " func_def
 " ------------------------------------------------------------------------
 function jedi#show_func_def()
-python << PYTHONEOF
-if 1:
-    row, column = vim.current.window.cursor
-    source = '\n'.join(vim.current.buffer)
-    buf_path = vim.current.buffer.name
-    call_def = functions.get_in_function_call(source, row, column, buf_path)
-    show_func_def(call_def)
-PYTHONEOF
+    python show_func_def(get_script().get_in_function_call())
     return ''
 endfunction
 
+
 function jedi#clear_func_def()
 python << PYTHONEOF
 if 1:
@@ -106,6 +103,7 @@ function! jedi#goto()
     python _goto()
 endfunction
 
+
 " ------------------------------------------------------------------------
 " get_definition
 " ------------------------------------------------------------------------
@@ -113,6 +111,7 @@ function! jedi#get_definition()
     python _goto(is_definition=True)
 endfunction
 
+
 " ------------------------------------------------------------------------
 " related_names
 " ------------------------------------------------------------------------
@@ -120,6 +119,7 @@ function! jedi#related_names()
     python _goto(is_related_name=True)
 endfunction
 
+
 " ------------------------------------------------------------------------
 " rename
 " ------------------------------------------------------------------------
@@ -143,7 +143,6 @@ if 1:
         # reset autocommand
         vim.command('autocmd! jedi_rename InsertLeave')
 
-        current_buf = vim.current.buffer.name
         replace = vim.eval("expand('<cword>')")
         vim.command('normal! u')  # undo new word
         vim.command('normal! u')  # 2u didn't work...
@@ -167,6 +166,7 @@ if 1:
 PYTHONEOF
 endfunction
 
+
 " ------------------------------------------------------------------------
 " show_pydoc
 " ------------------------------------------------------------------------
@@ -231,6 +231,7 @@ PYTHONEOF
     let b:current_syntax = "rst"
 endfunction
 
+
 " ------------------------------------------------------------------------
 " helper functions
 " ------------------------------------------------------------------------
@@ -245,6 +246,7 @@ function! jedi#new_buffer(path)
     endif
 endfunction
 
+
 function! jedi#tabnew(path)
 python << PYTHONEOF
 if 1:
@@ -271,6 +273,7 @@ if 1:
 PYTHONEOF
 endfunction
 
+
 function! s:add_goto_window()
     set lazyredraw
     cclose
@@ -283,6 +286,7 @@ function! s:add_goto_window()
     redraw!
 endfunction
 
+
 function! jedi#goto_window_on_enter()
     let l:list = getqflist()
     let l:data = l:list[line('.') - 1]
@@ -296,6 +300,7 @@ function! jedi#goto_window_on_enter()
     endif
 endfunction
 
+
 function! jedi#syn_stack()
     if !exists("*synstack")
         return []
@@ -303,6 +308,7 @@ function! jedi#syn_stack()
     return map(synstack(line('.'), col('.') - 1), 'synIDattr(v:val, "name")')
 endfunc
 
+
 function! jedi#do_popup_on_dot()
     let highlight_groups = jedi#syn_stack()
     for a in highlight_groups
@@ -435,21 +441,31 @@ class PythonToVimStr(str):
     def __repr__(self):
         return '"%s"' % self.replace('"', r'\"')
 
+
 def echo_highlight(msg):
     vim.command('echohl WarningMsg | echo "%s" | echohl None' % msg)
 
+
+def get_script(source=None, column=None):
+    if source is None:
+        source = '\n'.join(vim.current.buffer)
+    row = vim.current.window.cursor[0]
+    if column is None:
+        column = vim.current.window.cursor[1]
+    buf_path = vim.current.buffer.name
+    return functions.Script(source, row, column, buf_path)
+
+
 def _goto(is_definition=False, is_related_name=False, no_output=False):
     definitions = []
-    row, column = vim.current.window.cursor
-    buf_path = vim.current.buffer.name
-    source = '\n'.join(vim.current.buffer)
+    script = get_script()
     try:
         if is_related_name:
-            definitions = functions.related_names(source, row, column, buf_path)
+            definitions = script.related_names()
         elif is_definition:
-            definitions = functions.get_definition(source, row, column, buf_path)
+            definitions = script.get_definition()
         else:
-            definitions = functions.goto(source, row, column, buf_path)
+            definitions = script.goto()
     except functions.NotFoundError:
         echo_highlight("Cannot follow nothing. Put your cursor on a valid name.")
     except Exception:
@@ -512,7 +528,7 @@ def show_func_def(call_def, completion_lines=0):
     params = [p.get_code().replace('\n', '') for p in call_def.params]
     try:
         params[call_def.index] = '*%s*' % params[call_def.index]
-    except IndexError:
+    except (IndexError, TypeError):
         pass
 
     text = " (%s) " % ', '.join(params)
@@ -524,8 +540,6 @@ def show_func_def(call_def, completion_lines=0):
     repl = ("%s" + regex + "%s") % (line[:insert_column],
                     line[insert_column:end_column], text, line[end_column:])
     vim.eval('setline(%s, "%s")' % (row_to_replace, repl))
-    #vim.command(r"%ss/^.\{%s\}/\1%s/g" % (row_to_replace, column, text))
-
 PYTHONEOF
 
 " vim: set et ts=4:

test/regression.py

@@ -2,7 +2,9 @@
 import os
 import sys
 import unittest
+from os.path import abspath, dirname
 
+sys.path.append(abspath(dirname(abspath(__file__)) + '/..'))
 os.chdir(os.path.dirname(os.path.abspath(__file__)) + '/..')
 sys.path.append('.')
 
@@ -13,15 +15,18 @@ import functions
 
 class TestRegression(unittest.TestCase):
     def get_def(self, src, pos):
-        return functions.get_definition(src, pos[0], pos[1], '')
+        script = functions.Script(src, pos[0], pos[1], '')
+        return script.get_definition()
 
     def complete(self, src, pos):
-        return functions.complete(src, pos[0], pos[1], '')[0]
+        script = functions.Script(src, pos[0], pos[1], '')
+        return script.complete()
 
     def get_in_function_call(self, src, pos=None):
         if pos is None:
             pos = 1, len(src)
-        return functions.get_in_function_call(src, pos[0], pos[1], '')
+        script = functions.Script(src, pos[0], pos[1], '')
+        return script.get_in_function_call()
 
     def test_get_definition_cursor(self):
 
@@ -103,7 +108,7 @@ class TestRegression(unittest.TestCase):
         s = ("def abc(): pass\n"
              "abc.d.a.abc.d"
              )
-        functions.related_names(s, 2, 2, '/')
+        functions.Script(s, 2, 2, '/').related_names()
 
     def test_get_in_function_call(self):
         s = "isinstance(a, abs("

test/run.py

@@ -16,81 +16,37 @@ import debug
 sys.path.pop()  # pop again, because it might affect the completion
 
 
-def run_completion_test(correct, source, line_nr, index, line, path):
+def run_completion_test(script, correct, line_nr):
     """
     Runs tests for completions.
     Return if the test was a fail or not, with 1 for fail and 0 for success.
     """
-    # lines start with 1 and column is just the last (makes no
-    # difference for testing)
-    try:
-        completions, call_def = functions.complete(source, line_nr, index,
-                                                                        path)
-        #import cProfile as profile
-        #profile.run('functions.complete("""%s""", %i, %i, "%s")'
-        #                            % (source, line_nr, len(line), path))
-    except Exception:
-        print(traceback.format_exc())
-        print('test @%s: %s' % (line_nr - 1, line))
+    completions = script.complete()
+    #import cProfile; cProfile.run('script.complete()')
+
+    comp_str = set([c.word for c in completions])
+    if comp_str != set(literal_eval(correct)):
+        print('Solution @%s not right, received %s, wanted %s'\
+                    % (line_nr - 1, comp_str, correct))
         return 1
-    else:
-        # TODO remove set! duplicates should not be normal
-        comp_str = set([c.word for c in completions])
-        if comp_str != set(literal_eval(correct)):
-            print('Solution @%s not right, received %s, wanted %s'\
-                        % (line_nr - 1, comp_str, correct))
-            return 1
     return 0
 
 
-def run_definition_test(correct, source, line_nr, index, line, correct_start,
-                                                                        path):
+def run_definition_test(script, should_str, line_nr):
     """
     Runs tests for definitions.
     Return if the test was a fail or not, with 1 for fail and 0 for success.
     """
-    def defs(line_nr, indent):
-        return set(functions.get_definition(source, line_nr, indent, path))
-    try:
-        result = defs(line_nr, index)
-    except Exception:
-        print(traceback.format_exc())
-        print('test @%s: %s' % (line_nr - 1, line))
+    result = script.get_definition()
+    is_str = set(r.desc_with_module for r in result)
+    if is_str != should_str:
+        print('Solution @%s not right, received %s, wanted %s' \
+                    % (line_nr - 1, is_str, should_str))
         return 1
-    else:
-        should_be = set()
-        number = 0
-        for index in re.finditer('(?: +|$)', correct):
-            if correct == ' ':
-                continue
-            # -1 for the comment, +3 because of the comment start `#? `
-            start = index.start()
-            if print_debug:
-                functions.set_debug_function(None)
-            number += 1
-            try:
-                should_be |= defs(line_nr - 1, start + correct_start)
-            except Exception:
-                print(traceback.format_exc())
-                print('could not resolve %s indent %s' % (line_nr - 1, start))
-                return 1
-            if print_debug:
-                functions.set_debug_function(debug.print_to_stdout)
-        # because the objects have different ids, `repr` it, then compare it.
-        should_str = set(r.desc_with_module for r in should_be)
-        if len(should_str) < number:
-            print('Solution @%s not right, too few test results: %s' \
-                        % (line_nr - 1, should_str))
-            return 1
-        is_str = set(r.desc_with_module for r in result)
-        if is_str != should_str:
-            print('Solution @%s not right, received %s, wanted %s' \
-                        % (line_nr - 1, is_str, should_str))
-            return 1
     return 0
 
 
-def run_goto_test(correct, source, line_nr, index, line, path):
+def run_goto_test(script, correct, line_nr):
     """
     Runs tests for gotos.
     Tests look like this:
@@ -107,22 +63,16 @@ def run_goto_test(correct, source, line_nr, index, line, path):
 
     Return if the test was a fail or not, with 1 for fail and 0 for success.
     """
-    try:
-        result = functions.goto(source, line_nr, index, path)
-    except Exception:
-        print(traceback.format_exc())
-        print('test @%s: %s' % (line_nr - 1, line))
+    result = script.goto()
+    comp_str = str(sorted(r.description for r in result))
+    if comp_str != correct:
+        print('Solution @%s not right, received %s, wanted %s'\
+                    % (line_nr - 1, comp_str, correct))
         return 1
-    else:
-        comp_str = str(sorted(r.description for r in result))
-        if comp_str != correct:
-            print('Solution @%s not right, received %s, wanted %s'\
-                        % (line_nr - 1, comp_str, correct))
-            return 1
     return 0
 
 
-def run_related_name_test(correct, source, line_nr, index, line, path):
+def run_related_name_test(script, correct, line_nr):
     """
     Runs tests for gotos.
     Tests look like this:
@@ -132,20 +82,14 @@ def run_related_name_test(correct, source, line_nr, index, line, path):
 
     Return if the test was a fail or not, with 1 for fail and 0 for success.
     """
-    try:
-        result = functions.related_names(source, line_nr, index, path)
-    except Exception:
-        print(traceback.format_exc())
-        print('test @%s: %s' % (line_nr - 1, line))
+    result = script.related_names()
+    correct = correct.strip()
+    comp_str = set('(%s,%s)' % r.start_pos for r in result)
+    correct = set(correct.split(' ')) if correct else set()
+    if comp_str != correct:
+        print('Solution @%s not right, received %s, wanted %s'\
+                    % (line_nr - 1, comp_str, correct))
         return 1
-    else:
-        correct = correct.strip()
-        comp_str = set('(%s,%s)' % r.start_pos for r in result)
-        correct = set(correct.split(' ')) if correct else set()
-        if comp_str != correct:
-            print('Solution @%s not right, received %s, wanted %s'\
-                        % (line_nr - 1, comp_str, correct))
-            return 1
     return 0
 
 
@@ -164,9 +108,40 @@ def run_test(source, f_name, lines_to_execute):
     >>> #? int()
     >>> ab = 3; ab
     """
+    def get_defs(correct, correct_start, path):
+        def defs(line_nr, indent):
+            script = functions.Script(source, line_nr, indent, path)
+            return set(script.get_definition())
+
+        should_be = set()
+        number = 0
+        for index in re.finditer('(?: +|$)', correct):
+            if correct == ' ':
+                continue
+            # -1 for the comment, +3 because of the comment start `#? `
+            start = index.start()
+            if print_debug:
+                functions.set_debug_function(None)
+            number += 1
+            try:
+                should_be |= defs(line_nr - 1, start + correct_start)
+            except Exception:
+                raise Exception('could not resolve %s indent %s'
+                                                    % (line_nr - 1, start))
+            if print_debug:
+                functions.set_debug_function(debug.print_to_stdout)
+        # because the objects have different ids, `repr` it, then compare it.
+        should_str = set(r.desc_with_module for r in should_be)
+        if len(should_str) < number:
+            raise Exception('Solution @%s not right, too few test results: %s'
+                                                % (line_nr - 1, should_str))
+        return should_str
+
     fails = 0
     tests = 0
     correct = None
+    test_type = None
+    start = None
     for line_nr, line in enumerate(BytesIO(source.encode())):
         line = unicode(line)
         line_nr += 1
@@ -181,16 +156,21 @@ def run_test(source, f_name, lines_to_execute):
             # if a list is wanted, use the completion test, otherwise the
             # get_definition test
             path = completion_test_dir + os.path.sep + f_name
-            args = (correct, source, line_nr, index, line, path)
-            if test_type == '!':
-                fails += run_goto_test(*args)
-            elif test_type == '<':
-                fails += run_related_name_test(*args)
-            elif correct.startswith('['):
-                fails += run_completion_test(*args)
-            else:
-                fails += run_definition_test(correct, source, line_nr, index,
-                                                            line, start, path)
+            try:
+                script = functions.Script(source, line_nr, index, path)
+                if test_type == '!':
+                    fails += run_goto_test(script, correct, line_nr)
+                elif test_type == '<':
+                    fails += run_related_name_test(script, correct, line_nr)
+                elif correct.startswith('['):
+                    fails += run_completion_test(script, correct, line_nr)
+                else:
+                    should_str = get_defs(correct, start, path)
+                    fails += run_definition_test(script, should_str, line_nr)
+            except Exception:
+                print(traceback.format_exc())
+                print('test @%s: %s' % (line_nr - 1, line))
+                fails += 1
             correct = None
             tests += 1
         else: