Return types for docstrings

jedi/docstrings.py

@@ -11,6 +11,16 @@ DOCSTRING_PARAM_PATTERNS = [
     r'\s*%s\s+\(([^()]+)\)'         # googley
 ]
 
+DOCSTRING_RETURN_PATTERNS = [
+        re.compile(r'\s*:rtype:\s*([^\n]+)', re.M), # Sphinx
+        re.compile(r'\s*@rtype:\s*([^\n]+)', re.M), # Epidoc
+
+        # Googley is the most undocumented format for
+        # return types, so we try to analyze the first
+        # line or line after `Returns:` keyword
+        re.compile(r'returns\s*:[\s\n]*([^\n]+)', re.M|re.I),
+]
+
 #@cache.memoize_default()  # TODO add
 def follow_param(param):
     func = param.parent_function
@@ -35,3 +45,25 @@ def search_param_in_docstr(docstr, param_str):
                 return match.group(1)
 
     return None
+
+
+def find_return_types(func):
+    if isinstance(func, evaluate.InstanceElement):
+        func = func.var
+
+    if isinstance(func, evaluate.Function):
+        func = func.base_func
+
+    type_str = search_return_in_docstr(func.docstr)
+    if not type_str:
+        return []
+
+    p = parsing.PyFuzzyParser(type_str, None, (1, 0), no_docstr=True)
+    p.user_stmt.parent = func
+    return list(evaluate.follow_statement(p.user_stmt))
+
+def search_return_in_docstr(code):
+    for p in DOCSTRING_RETURN_PATTERNS:
+        match = p.search(code)
+        if match:
+            return match.group(1)

jedi/evaluate.py

@@ -472,7 +472,7 @@ class Execution(Executable):
         if func.is_generator and not evaluate_generator:
             return [Generator(func, self.var_args)]
         else:
-            stmts = []
+            stmts = docstrings.find_return_types(func)
             for r in self.returns:
                 if r is not None:
                     stmts += follow_statement(r)

test/completion/docstring.py

@@ -4,28 +4,79 @@ def f(a, b):
     """ asdfasdf
     :param a: blablabla
     :type a: str
+    :rtype: dict
     """
     #? str()
     a
     #? 
     b
 
+#? dict()
+f()
+
 def g(a, b):
     """ asdfasdf
     Arguments:
         a (str): blablabla
+
+    Returns: list
+        Blah blah.
     """
     #? str()
     a
     #? 
     b
 
+#? list()
+g()
+
 def e(a, b):
     """ asdfasdf
     @type a: str
     @param a: blablabla
+    @rtype: list
     """
     #? str()
     a
     #?
     b
+
+#? list()
+e()
+
+
+# Returns with param type only
+def rparam(a,b):
+    """
+    @type a: str
+    """
+    return a
+
+#? str()
+rparam()
+
+
+# Composite types
+def composite():
+    """
+    @rtype: (str, int, dict)
+    """
+
+x, y, z = composite()
+#? str()
+x
+#? int()
+y
+#? dict()
+z
+
+
+# Both docstring and calculated return type
+def both():
+    """
+    @rtype: str
+    """
+    return 23
+
+#? str(), int()
+both()