docstring parsing to infer argument types, fixes #40

jedi/docstrings.py

@@ -0,0 +1,34 @@
+""" Processing of docstrings, which means parsing for types. """
+
+import re
+
+import evaluate
+
+
+#@evaluate.memoize_default()  # TODO add
+def follow_param(param):
+    func = param.parent_function()
+    #print func, param, param.parent_function()
+    param_str = search_param_in_docstr(func.docstr, str(param.get_name()))
+
+    if param_str is not None:
+        scope = func.get_parent_until()
+        return evaluate.get_scopes_for_name(scope, param_str,
+                                            search_global=True)
+    return []
+
+
+def search_param_in_docstr(docstr, param_str):
+    lines = docstr.split('\n')
+
+    # look at #40 to see definitions of those params 
+    sphinx_comp = ':type %s:' % param_str
+    googley_comp = re.compile('\s*%s\s+\(([^()]+)\)' % re.escape(param_str))
+    for l in lines:
+        if l.startswith(sphinx_comp):
+            return l.replace(sphinx_comp, '', 1).strip()
+
+        r = re.match(googley_comp, l)
+        if r is not None:
+            return r.group(1)
+    return None

jedi/evaluate.py

@@ -25,6 +25,7 @@ import builtin
 import imports
 import helpers
 import dynamic
+import docstrings
 
 memoize_caches = []
 faked_scopes = []
@@ -1064,6 +1065,12 @@ def get_scopes_for_name(scope, name_str, position=None, search_global=False,
                                 and r.position_nr > 0:  # 0 would be self
                             r = func.var.params[r.position_nr]
 
+                        # add docstring knowledge
+                        doc_params = docstrings.follow_param(r)
+                        if doc_params:
+                            res_new += doc_params
+                            continue
+
                         if not r.is_generated:
                             res_new += dynamic.search_params(r)
                             if not r.assignment_details:

jedi/parsing.py

@@ -369,6 +369,7 @@ class Function(Scope):
         self.params = params
         for p in params:
             p.parent = weakref.ref(self)
+            p.parent_function = weakref.ref(self)
         self.decorators = []
         self.returns = []
         self.is_generator = False
@@ -872,6 +873,7 @@ class Param(Statement):
         self.position_nr = None
         self.is_generated = False
         self.annotation_stmt = None
+        self.parent_function = None
 
     def add_annotation(self, annotation_stmt):
         annotation_stmt.parent = weakref.ref(self)