Fix instance docstring

davidhalter · Jul 12, 2017 · f7fbf41 · f7fbf41
1 parent 9fb7fb6
commit f7fbf41
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/jedi/evaluate/instance.py b/jedi/evaluate/instance.py
@@ -24,6 +24,7 @@ def __init__(self, evaluator, parent_context, class_context, var_args):
         super(AbstractInstanceContext, self).__init__(evaluator, parent_context)
         # Generated instances are classes that are just generated by self
         # (No var_args) used.
+        self.tree_node = class_context.tree_node
         self.class_context = class_context
         self.var_args = var_args
 

diff --git a/test/test_evaluate/test_docstring.py b/test/test_evaluate/test_docstring.py
@@ -22,6 +22,21 @@ def func():
         func""").goto_definitions()
         self.assertEqual(defs[0].raw_doc, 'Docstring of `func`.')
 
+    def test_class_doc(self):
+        defs = jedi.Script("""
+        class TestClass():
+            '''Docstring of `TestClass`.'''
+        TestClass""").goto_definitions()
+        self.assertEqual(defs[0].docstring(), 'Docstring of `TestClass`.')
+
+    def test_instance_doc(self):
+        defs = jedi.Script("""
+        class TestClass():
+            '''Docstring of `TestClass`.'''
+        tc = TestClass()
+        tc""").goto_definitions()
+        self.assertEqual(defs[0].docstring(), 'Docstring of `TestClass`.')
+
     @unittest.skip('need evaluator class for that')
     def test_attribute_docstring(self):
         defs = jedi.Script("""