Better control over docstring generation

jedi/evaluate/signature.py

@@ -41,6 +41,8 @@ class TreeSignature(AbstractSignature):
         return get_call_signature(
             self._function_context.tree_node,
             call_string=self.name.string_name,
+            omit_first_param=self.is_bound,
+            omit_return_annotation=self.context.is_class(),
         )
 
 

jedi/parser_utils.py

@@ -138,7 +138,8 @@ def safe_literal_eval(value):
         return ''
 
 
-def get_call_signature(funcdef, width=72, call_string=None):
+def get_call_signature(funcdef, width=72, call_string=None,
+                       omit_first_param=False, omit_return_annotation=False):
     """
     Generate call signature of this function.
 
@@ -155,12 +156,14 @@ def get_call_signature(funcdef, width=72, call_string=None):
             call_string = '<lambda>'
         else:
             call_string = funcdef.name.value
-    if funcdef.type == 'lambdef':
-        p = '(' + ''.join(param.get_code() for param in funcdef.get_params()).strip() + ')'
-    else:
-        p = funcdef.children[2].get_code()
+    params = funcdef.get_params()
+    if omit_first_param:
+        params = params[1:]
+    print(params, omit_first_param)
+    p = '(' + ''.join(param.get_code() for param in params).strip() + ')'
+    # TODO this is pretty bad, we should probably just normalize.
     p = re.sub(r'\s+', ' ', p)
-    if funcdef.annotation:
+    if funcdef.annotation and not omit_return_annotation:
         rtype = " ->" + funcdef.annotation.get_code()
     else:
         rtype = ""

test/test_evaluate/test_docstring.py

@@ -35,7 +35,20 @@ def test_class_doc(Script):
     class TestClass():
         '''Docstring of `TestClass`.'''
     TestClass""").goto_definitions()
-    assert defs[0].docstring() == 'Docstring of `TestClass`.'
+
+    expected = 'Docstring of `TestClass`.'
+    assert defs[0].docstring(raw=True) == expected
+    assert defs[0].docstring() == 'TestClass()\n\n' + expected
+
+
+def test_class_doc_with_init(Script):
+    d, = Script("""
+    class TestClass():
+        '''Docstring'''
+        def __init__(self, foo, bar=3): pass
+    TestClass""").goto_definitions()
+
+    assert d.docstring() == 'TestClass(foo, bar=3)\n\nDocstring'
 
 
 def test_instance_doc(Script):