From 45a5eee18ae3e025f671c89d81fbe65299029e68 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Sun, 5 May 2019 19:50:52 +0200
Subject: [PATCH] Better control over docstring generation

---
 jedi/evaluate/signature.py           |  2 ++
 jedi/parser_utils.py                 | 15 +++++++++------
 test/test_evaluate/test_docstring.py | 15 ++++++++++++++-
 3 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/jedi/evaluate/signature.py b/jedi/evaluate/signature.py
index 8150294f..9b3fc2df 100644
--- a/jedi/evaluate/signature.py
+++ b/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(),
         )
 
 
diff --git a/jedi/parser_utils.py b/jedi/parser_utils.py
index d4f8143a..3fc3019d 100644
--- a/jedi/parser_utils.py
+++ b/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 = ""
diff --git a/test/test_evaluate/test_docstring.py b/test/test_evaluate/test_docstring.py
index 649d1ce3..676f4e8f 100644
--- a/test/test_evaluate/test_docstring.py
+++ b/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):