Add basic numpydoc support

jedi/evaluate/docstrings.py

@@ -14,6 +14,7 @@ As an addition to parameter searching, this module also provides return
 annotations.
 """
 
+from ast import literal_eval
 import re
 from itertools import chain
 from textwrap import dedent
@@ -36,28 +37,44 @@ DOCSTRING_RETURN_PATTERNS = [
 REST_ROLE_PATTERN = re.compile(r':[^`]+:`([^`]+)`')
 
 
-@memoize_default(None, evaluator_is_first_arg=True)
-def follow_param(evaluator, param):
-    func = param.parent_function
-    param_str = _search_param_in_docstr(func.raw_doc, str(param.get_name()))
-    return _evaluate_for_statement_string(evaluator, param_str, param.get_parent_until())
+try:
+    from numpydoc.docscrape import NumpyDocString
+except ImportError:
+    def _search_param_in_numpydocstr(docstr, param_str):
+        return []
+else:
+    def _search_param_in_numpydocstr(docstr, param_str):
+        """Search `docstr` (in numpydoc format) for type(-s) of `param_str`."""
+        params = NumpyDocString(docstr)._parsed_data['Parameters']
+        for p_name, p_type, p_descr in params:
+            if p_name == param_str:
+                m = re.match('([^,]+(,[^,]+)*?)(,[ ]*optional)?$', p_type)
+                if m:
+                    p_type = m.group(1)
+
+                if p_type.startswith('{'):
+                    types = set(type(x).__name__ for x in literal_eval(p_type))
+                    return list(types)
+                else:
+                    return [p_type]
+        return []
 
 
 def _search_param_in_docstr(docstr, param_str):
     """
-    Search `docstr` for a type of `param_str`.
+    Search `docstr` for type(-s) of `param_str`.
 
     >>> _search_param_in_docstr(':type param: int', 'param')
-    'int'
+    ['int']
     >>> _search_param_in_docstr('@type param: int', 'param')
-    'int'
+    ['int']
     >>> _search_param_in_docstr(
     ...   ':type param: :class:`threading.Thread`', 'param')
-    'threading.Thread'
-    >>> _search_param_in_docstr('no document', 'param') is None
-    True
+    ['threading.Thread']
+    >>> bool(_search_param_in_docstr('no document', 'param'))
+    False
     >>> _search_param_in_docstr(':param int param: some description', 'param')
-    'int'
+    ['int']
 
     """
     # look at #40 to see definitions of those params
@@ -66,9 +83,10 @@ def _search_param_in_docstr(docstr, param_str):
     for pattern in patterns:
         match = pattern.search(docstr)
         if match:
-            return _strip_rst_role(match.group(1))
+            return [_strip_rst_role(match.group(1))]
 
-    return None
+    return (_search_param_in_numpydocstr(docstr, param_str) or
+            [])
 
 
 def _strip_rst_role(type_str):
@@ -129,6 +147,17 @@ def _evaluate_for_statement_string(evaluator, string, module):
     return list(chain.from_iterable(it)) or definitions
 
 
+@memoize_default(None, evaluator_is_first_arg=True)
+def follow_param(evaluator, param):
+    func = param.parent_function
+
+    return [p
+            for param_str in _search_param_in_docstr(func.raw_doc,
+                                                     str(param.get_name()))
+            for p in _evaluate_for_statement_string(evaluator, param_str,
+                                                    param.get_parent_until())]
+
+
 @memoize_default(None, evaluator_is_first_arg=True)
 def find_return_types(evaluator, func):
     def search_return_in_docstr(code):

test/test_evaluate/test_docstring.py

@@ -6,6 +6,13 @@ from textwrap import dedent
 import jedi
 from ..helpers import unittest
 
+try:
+    import numpydoc
+except ImportError:
+    numpydoc_unavailable = True
+else:
+    numpydoc_unavailable = False
+
 
 class TestDocstring(unittest.TestCase):
     def test_function_doc(self):
@@ -97,3 +104,49 @@ class TestDocstring(unittest.TestCase):
         assert 'a' in names
         assert '__init__' in names
         assert 'mro' not in names  # Exists only for types.
+
+    @unittest.skipIf(numpydoc_unavailable, 'numpydoc module is unavailable')
+    def test_numpydoc_docstring(self):
+        s = dedent('''
+        def foobar(x, y):
+            """
+            Parameters
+            ----------
+            x : int
+            y : str
+            """
+            y.''')
+        names = [c.name for c in jedi.Script(s).completions()]
+        assert 'isupper' in names
+        assert 'capitalize' in names
+
+    @unittest.skipIf(numpydoc_unavailable, 'numpydoc module is unavailable')
+    def test_numpydoc_docstring_set_of_values(self):
+        s = dedent('''
+        def foobar(x, y):
+            """
+            Parameters
+            ----------
+            x : {'foo', 'bar', 100500}, optional
+            """
+            x.''')
+        names = [c.name for c in jedi.Script(s).completions()]
+        assert 'isupper' in names
+        assert 'capitalize' in names
+        assert 'numerator' in names
+
+    @unittest.skipIf(numpydoc_unavailable, 'numpydoc module is unavailable')
+    def test_numpydoc_alternative_types(self):
+        s = dedent('''
+        def foobar(x, y):
+            """
+            Parameters
+            ----------
+            x : int or str or list
+            """
+            x.''')
+        names = [c.name for c in jedi.Script(s).completions()]
+        assert 'isupper' in names
+        assert 'capitalize' in names
+        assert 'numerator' in names
+        assert 'append' in names