Return types from numpydocs and compiled objects return types from docstrings

jedi/evaluate/compiled/__init__.py

@@ -205,9 +205,9 @@ def name(self):
         return CompiledContextName(self, name)
 
     def _execute_function(self, params):
+        from jedi.evaluate import docstrings
         if self.type != 'funcdef':
             return
-
         for name in self._parse_function_doc()[1].split():
             try:
                 bltn_obj = getattr(_builtins, name)
@@ -221,6 +221,8 @@ def _execute_function(self, params):
                 bltn_obj = create(self.evaluator, bltn_obj)
                 for result in self.evaluator.execute(bltn_obj, params):
                     yield result
+        for type_ in docstrings.infer_return_types(self):
+            yield type_
 
     def get_self_attributes(self):
         return []  # Instance compatibility

jedi/evaluate/docstrings.py

@@ -1,11 +1,12 @@
 """
 Docstrings are another source of information for functions and classes.
 :mod:`jedi.evaluate.dynamic` tries to find all executions of functions, while
-the docstring parsing is much easier. There are two different types of
+the docstring parsing is much easier. There are three different types of
 docstrings that |jedi| understands:
 
 - `Sphinx <http://sphinx-doc.org/markup/desc.html#info-field-lists>`_
 - `Epydoc <http://epydoc.sourceforge.net/manual-fields.html>`_
+- `Numpydoc <https://github.com/numpy/numpy/blob/master/doc/HOWTO_DOCUMENT.rst.txt>`_
 
 For example, the sphinx annotation ``:type foo: str`` clearly states that the
 type of ``foo`` is ``str``.
@@ -46,23 +47,67 @@
 except ImportError:
     def _search_param_in_numpydocstr(docstr, param_str):
         return []
+
+    def _search_return_in_numpydocstr(docstr):
+        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']
+        try:
+            # This is a non-public API. If it ever changes we should be 
+            # prepared and return gracefully.
+            params = NumpyDocString(docstr)._parsed_data['Parameters']
+        except (KeyError, AttributeError):
+            return []
         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 _expand_typestr(p_type)
         return []
 
+    def _search_return_in_numpydocstr(docstr):
+        """
+        Search `docstr` (in numpydoc format) for type(-s) of function returns.
+        """
+        doc = NumpyDocString(docstr)
+        try:
+            # This is a non-public API. If it ever changes we should be 
+            # prepared and return gracefully.
+            returns = doc._parsed_data['Returns']
+            returns += doc._parsed_data['Yields']
+        except (KeyError, AttributeError):
+            raise StopIteration
+        for r_name, r_type, r_descr in returns:
+            #Return names are optional and if so the type is in the name
+            if not r_type:
+                r_type = r_name
+            for type_ in _expand_typestr(r_type):
+                yield type_
+
+
+def _expand_typestr(type_str):
+    """
+    Attempts to interpret the possible types in `type_str`
+    """
+    # Check if alternative types are specified with 'or'
+    if re.search('\\bor\\b', type_str):
+        types = [t.split('of')[0].strip() for t in type_str.split('or')]
+    # Check if like "list of `type`" and set type to list
+    elif re.search('\\bof\\b', type_str):
+        types = [type_str.split('of')[0]]
+    # Check if type has is a set of valid literal values eg: {'C', 'F', 'A'}
+    elif type_str.startswith('{'):
+        # python2 does not support literal set evals
+        # workaround this by using lists instead
+        type_str = type_str.replace('{', '[').replace('}', ']')
+        types = set(type(x).__name__ for x in literal_eval(type_str))
+    # Otherwise just return the typestr wrapped in a list
+    else:
+        types = [type_str]
+    return types
+
 
 def _search_param_in_docstr(docstr, param_str):
     """
@@ -213,7 +258,12 @@ def search_return_in_docstr(code):
         for p in DOCSTRING_RETURN_PATTERNS:
             match = p.search(code)
             if match:
-                return _strip_rst_role(match.group(1))
+                yield _strip_rst_role(match.group(1))
+        # Check for numpy style return hint
+        for type_ in _search_return_in_numpydocstr(code):
+            yield type_
+
+    for type_str in search_return_in_docstr(function_context.py__doc__()):
+        for type_eval in _evaluate_for_statement_string(function_context.get_root_context(), type_str):
+            yield type_eval
 
-    type_str = search_return_in_docstr(function_context.py__doc__())
-    return _evaluate_for_statement_string(function_context.get_root_context(), type_str)

test/test_evaluate/test_docstring.py

@@ -4,14 +4,22 @@
 
 from textwrap import dedent
 import jedi
+import pytest
 from ..helpers import unittest
 
 try:
-    import numpydoc
+    import numpydoc  # NOQA
 except ImportError:
     numpydoc_unavailable = True
 else:
     numpydoc_unavailable = False
+
+try:
+    import numpy
+except ImportError:
+    numpy_unavailable = True
+else:
+    numpy_unavailable = False
 
 
 class TestDocstring(unittest.TestCase):
@@ -124,48 +132,200 @@ def test_docstring_keyword(self):
         completions = jedi.Script('assert').completions()
         self.assertIn('assert', completions[0].docstring())
 
-    @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
+# ---- Numpy Style Tests ---
+
+@pytest.mark.skipif(numpydoc_unavailable, 
+                    reason='numpydoc module is unavailable')
+def test_numpydoc_parameters():
+    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
+
+@pytest.mark.skipif(numpydoc_unavailable, 
+                    reason='numpydoc module is unavailable')
+def test_numpydoc_parameters_set_of_values():
+    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
+
+@pytest.mark.skipif(numpydoc_unavailable, 
+                    reason='numpydoc module is unavailable')
+def test_numpydoc_parameters_alternative_types():
+    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
+
+@pytest.mark.skipif(numpydoc_unavailable, 
+                    reason='numpydoc module is unavailable')
+def test_numpydoc_returns():
+    s = dedent('''
+    def foobar():
+        """
+        Returns
+        ----------
+        x : int
+        y : str
+        """
+        return x
+
+    def bazbiz():
+        z = foobar()
+        z.''')
+    names = [c.name for c in jedi.Script(s).completions()]
+    assert 'isupper' in names
+    assert 'capitalize' in names
+    assert 'numerator' in names
+
+@pytest.mark.skipif(numpydoc_unavailable, 
+                    reason='numpydoc module is unavailable')
+def test_numpydoc_returns_set_of_values():
+    s = dedent('''
+    def foobar():
+        """
+        Returns
+        ----------
+        x : {'foo', 'bar', 100500}
+        """
+        return x
+
+    def bazbiz():
+        z = foobar()
+        z.''')
+    names = [c.name for c in jedi.Script(s).completions()]
+    assert 'isupper' in names
+    assert 'capitalize' in names
+    assert 'numerator' in names
+
+@pytest.mark.skipif(numpydoc_unavailable, 
+                    reason='numpydoc module is unavailable')
+def test_numpydoc_returns_alternative_types():
+    s = dedent('''
+    def foobar():
+        """
+        Returns
+        ----------
+        int or list of str
+        """
+        return x
+
+    def bazbiz():
+        z = foobar()
+        z.''')
+    names = [c.name for c in jedi.Script(s).completions()]
+    assert 'isupper' not in names
+    assert 'capitalize' not in names
+    assert 'numerator' in names
+    assert 'append' in names
+
+@pytest.mark.skipif(numpydoc_unavailable, 
+                    reason='numpydoc module is unavailable')
+def test_numpydoc_returns_list_of():
+    s = dedent('''
+    def foobar():
+        """
+        Returns
+        ----------
+        list of str
+        """
+        return x
+
+    def bazbiz():
+        z = foobar()
+        z.''')
+    names = [c.name for c in jedi.Script(s).completions()]
+    assert 'append' in names
+    assert 'isupper' not in names
+    assert 'capitalize' not in names
+
+@pytest.mark.skipif(numpydoc_unavailable, 
+                    reason='numpydoc module is unavailable')
+def test_numpydoc_returns_obj():
+    s = dedent('''
+    def foobar(x, y):
+        """
+        Returns
+        ----------
+        int or random.Random
+        """
+        return x + y
+
+    def bazbiz():
+        z = foobar(x, y)
+        z.''')
+    script = jedi.Script(s)
+    names = [c.name for c in script.completions()]
+    assert 'numerator' in names
+    assert 'seed' in names
+
+@pytest.mark.skipif(numpydoc_unavailable, 
+                    reason='numpydoc module is unavailable')
+def test_numpydoc_yields():
+    s = dedent('''
+    def foobar():
+        """
+        Yields
+        ----------
+        x : int
+        y : str
+        """
+        return x
+
+    def bazbiz():
+        z = foobar():
+        z.''')
+    names = [c.name for c in jedi.Script(s).completions()]
+    print('names',names)
+    assert 'isupper' in names
+    assert 'capitalize' in names
+    assert 'numerator' in names
+
+@pytest.mark.skipif(numpydoc_unavailable or numpy_unavailable, 
+                    reason='numpydoc or numpy module is unavailable')
+def test_numpy_returns():
+    s = dedent('''
+    import numpy
+    x = numpy.asarray([])
+    x.d''')
+    names = [c.name for c in jedi.Script(s).completions()]
+    print(names)
+    assert 'diagonal' in names
+
+@pytest.mark.skipif(numpydoc_unavailable or numpy_unavailable, 
+                    reason='numpydoc or numpy module is unavailable')
+def test_numpy_comp_returns():
+    s = dedent('''
+    import numpy
+    x = numpy.array([])
+    x.d''')
+    names = [c.name for c in jedi.Script(s).completions()]
+    print(names)
+    assert 'diagonal' in names
+