From e086c433ffdb05fc00f80c6cef50c58cc4d2c46b Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Tue, 4 Sep 2018 10:08:09 +0200
Subject: [PATCH] Fix compiled docstrings for stubs

---
 jedi/parser_utils.py                | 22 +++++++++++++---------
 jedi/plugins/typeshed.py            | 22 ++++++++++++++++++++--
 test/test_evaluate/test_compiled.py |  4 +++-
 3 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/jedi/parser_utils.py b/jedi/parser_utils.py
index bb033ed0..ea15dff9 100644
--- a/jedi/parser_utils.py
+++ b/jedi/parser_utils.py
@@ -167,19 +167,23 @@ def get_call_signature(funcdef, width=72, call_string=None):
     return '\n'.join(textwrap.wrap(code, width))
 
 
+def get_call_signature_for_any(any_node):
+    call_signature = None
+    if any_node.type == 'classdef':
+        for funcdef in any_node.iter_funcdefs():
+            if funcdef.name.value == '__init__':
+                call_signature = \
+                    get_call_signature(funcdef, call_string=any_node.name.value)
+    elif any_node.type in ('funcdef', 'lambdef'):
+        call_signature = get_call_signature(any_node)
+    return call_signature
+
+
 def get_doc_with_call_signature(scope_node):
     """
     Return a document string including call signature.
     """
-    call_signature = None
-    if scope_node.type == 'classdef':
-        for funcdef in scope_node.iter_funcdefs():
-            if funcdef.name.value == '__init__':
-                call_signature = \
-                    get_call_signature(funcdef, call_string=scope_node.name.value)
-    elif scope_node.type in ('funcdef', 'lambdef'):
-        call_signature = get_call_signature(scope_node)
-
+    call_signature = get_call_signature_for_any(scope_node)
     doc = clean_scope_docstring(scope_node)
     if call_signature is None:
         return doc
diff --git a/jedi/plugins/typeshed.py b/jedi/plugins/typeshed.py
index e1a45055..4ec3b63d 100644
--- a/jedi/plugins/typeshed.py
+++ b/jedi/plugins/typeshed.py
@@ -6,7 +6,9 @@ from jedi._compatibility import FileNotFoundError
 from jedi.plugins.base import BasePlugin
 from jedi.evaluate.cache import evaluator_function_cache
 from jedi.cache import memoize_method
-from jedi.evaluate.base_context import ContextSet, iterator_to_context_set
+from jedi.parser_utils import get_call_signature_for_any
+from jedi.evaluate.base_context import ContextSet, iterator_to_context_set, \
+    ContextWrapper
 from jedi.evaluate.filters import AbstractTreeName, ParserTreeFilter, \
     TreeNameDefinition, NameWrapper
 from jedi.evaluate.context import ModuleContext, FunctionContext, \
@@ -187,7 +189,9 @@ class NameWithStubMixin(object):
         # This basically merges stub contexts with actual contexts.
         for actual_context in actual_contexts:
             for stub_context in stub_contexts:
-                if isinstance(stub_context, FunctionContext) \
+                if isinstance(actual_context, CompiledObject):
+                    yield StubContextWithCompiled(stub_context, actual_context)
+                elif isinstance(stub_context, FunctionContext) \
                         and isinstance(actual_context, FunctionContext):
                     yield StubFunctionContext(
                         actual_context.evaluator,
@@ -380,6 +384,20 @@ class StubOnlyModuleContext(ModuleContext):
             yield f
 
 
+class StubContextWithCompiled(ContextWrapper):
+    def __init__(self, stub_context, compiled_context):
+        super(StubContextWithCompiled, self).__init__(stub_context)
+        self._compiled_context = compiled_context
+
+    def py__doc__(self, include_call_signature=False):
+        doc = self._compiled_context.py__doc__()
+        if include_call_signature:
+            call_sig = get_call_signature_for_any(self._wrapped_context.tree_node)
+            if call_sig is not None:
+                doc = call_sig + '\n\n' + doc
+        return doc
+
+
 class TypingModuleWrapper(StubOnlyModuleContext):
     def get_filters(self, *args, **kwargs):
         filters = super(TypingModuleWrapper, self).get_filters(*args, **kwargs)
diff --git a/test/test_evaluate/test_compiled.py b/test/test_evaluate/test_compiled.py
index 1f64f004..38d89a0d 100644
--- a/test/test_evaluate/test_compiled.py
+++ b/test/test_evaluate/test_compiled.py
@@ -77,7 +77,9 @@ def test_method_completion(Script, environment):
 def test_time_docstring(Script):
     import time
     comp, = Script('import time\ntime.sleep').completions()
-    assert comp.docstring() == time.sleep.__doc__
+    assert comp.docstring(raw=True) == time.sleep.__doc__
+    expected = 'sleep(secs: float) -> None\n\n' + time.sleep.__doc__
+    assert comp.docstring() == expected
 
 
 def test_dict_values(Script):