Merge pull request #141 from tkf/attribute-docstring

Attribute docstring support

jedi/evaluate.py

@@ -259,6 +259,9 @@ def find_name(scope, name_str, position=None, search_global=False,
                                 # so just ignore it.
                                 continue
 
+                    if r.docstr:
+                        res_new.append(r)
+
                     scopes = follow_statement(r, seek_name=name_str)
                     add += remove_statements(scopes)
 

jedi/parsing.py

@@ -395,6 +395,19 @@ class Parser(object):
 
             self._check_user_stmt(stmt)
 
+        # Attribute docstring (PEP 257) support
+        try:
+            # If string literal is being parsed
+            first_tok = stmt.token_list[0]
+            if (not stmt.set_vars and
+                not stmt.used_vars and
+                len(stmt.token_list) == 1 and
+                first_tok[0] == tokenize.STRING):
+                # ... then set it as a docstring
+                self.scope.statements[-1].add_docstr(first_tok[1])
+        except (IndexError, AttributeError):
+            pass
+
         if tok in always_break + not_first_break:
             self._gen.push_last_back()
         return stmt, tok

jedi/parsing_representation.py

@@ -733,7 +733,8 @@ class Statement(Simple):
     :param  start_pos: Position (line, column) of the Statement.
     """
     __slots__ = ('token_list', 'used_vars',
-                 'set_vars', '_commands', '_assignment_details')
+                 'set_vars', '_commands', '_assignment_details',
+                 'docstr')
 
     def __init__(self, module, set_vars, used_vars, token_list,
                  start_pos, end_pos, parent=None):
@@ -744,12 +745,17 @@ class Statement(Simple):
             s.parent = self.use_as_parent
         self.set_vars = self._remove_executions_from_set_vars(set_vars)
         self.parent = parent
+        self.docstr = ''
 
         # cache
         self._commands = None
         self._assignment_details = []
         # this is important for other scripts
 
+    def add_docstr(self, string):
+        """ Clean up a docstring """
+        self.docstr = cleandoc(literal_eval(string))
+
     def _remove_executions_from_set_vars(self, set_vars):
         """
         Important mainly for assosiative arrays::
@@ -1339,6 +1345,11 @@ class Name(Simple):
         """ Returns the names in a full string format """
         return ".".join(self.names)
 
+    @property
+    def docstr(self):
+        """Return attribute docstring (PEP 257) if exists."""
+        return self.parent.docstr
+
     def __str__(self):
         return self.get_code()
 

test/regression.py

@@ -352,6 +352,33 @@ class TestRegression(TestBase):
         assert 'start' in words
 
 
+class TestDocstring(TestBase):
+
+    def test_function_doc(self):
+        defs = self.definition("""
+        def func():
+            '''Docstring of `func`.'''
+        func""")
+        self.assertEqual(defs[0].raw_doc, 'Docstring of `func`.')
+
+    def test_attribute_docstring(self):
+        defs = self.definition("""
+        x = None
+        '''Docstring of `x`.'''
+        x""")
+        self.assertEqual(defs[0].raw_doc, 'Docstring of `x`.')
+
+    def test_multiple_docstrings(self):
+        defs = self.definition("""
+        def func():
+            '''Original docstring.'''
+        x = func
+        '''Docstring of `x`.'''
+        x""")
+        docs = [d.raw_doc for d in defs]
+        self.assertEqual(docs, ['Original docstring.', 'Docstring of `x`.'])
+
+
 class TestFeature(TestBase):
     def test_full_name(self):
         """ feature request #61"""