dostring fixes.

jedi/evaluate/docstrings.py

@@ -16,12 +16,14 @@ annotations.
 
 from ast import literal_eval
 import re
-from itertools import chain
 from textwrap import dedent
 
+from jedi._compatibility import u
+from jedi.common import unite
+from jedi.evaluate import context
 from jedi.evaluate.cache import memoize_default
 from jedi.parser import ParserWithRecovery, load_grammar
-from jedi.parser.tree import Class
+from jedi.parser.tree import search_ancestor
 from jedi.common import indent_block
 from jedi.evaluate.iterable import ArrayLiteralContext, FakeSequence, AlreadyEvaluated
 
@@ -114,12 +116,12 @@ def _strip_rst_role(type_str):
         return type_str
 
 
-def _evaluate_for_statement_string(evaluator, string, module):
-    code = dedent("""
+def _evaluate_for_statement_string(module_context, string):
+    code = dedent(u("""
     def pseudo_docstring_stuff():
         # Create a pseudo function for docstring statements.
-    %s
-    """)
+    {0}
+    """))
     if string is None:
         return []
 
@@ -131,31 +133,39 @@ def _evaluate_for_statement_string(evaluator, string, module):
     # Take the default grammar here, if we load the Python 2.7 grammar here, it
     # will be impossible to use `...` (Ellipsis) as a token. Docstring types
     # don't need to conform with the current grammar.
-    p = ParserWithRecovery(load_grammar(), code % indent_block(string))
+    p = ParserWithRecovery(load_grammar(), code.format(indent_block(string)))
     try:
-        pseudo_cls = p.module.subscopes[0]
+        funcdef = p.module.subscopes[0]
         # First pick suite, then simple_stmt and then the node,
         # which is also not the last item, because there's a newline.
-        stmt = pseudo_cls.children[-1].children[-1].children[-2]
+        stmt = funcdef.children[-1].children[-1].children[-2]
     except (AttributeError, IndexError):
         return []
 
+    from jedi.evaluate.param import ValuesArguments
+    from jedi.evaluate.representation import FunctionExecutionContext
+    func_context = FunctionExecutionContext(
+        module_context.evaluator,
+        module_context,
+        funcdef,
+        ValuesArguments([])
+    )
+
     # Use the module of the param.
     # TODO this module is not the module of the param in case of a function
     # call. In that case it's the module of the function call.
     # stuffed with content from a function call.
-    pseudo_cls.parent = module
-    return list(_execute_types_in_stmt(evaluator, stmt))
+    return list(_execute_types_in_stmt(func_context, stmt))
 
 
-def _execute_types_in_stmt(evaluator, stmt):
+def _execute_types_in_stmt(module_context, stmt):
     """
     Executing all types or general elements that we find in a statement. This
     doesn't include tuple, list and dict literals, because the stuff they
     contain is executed. (Used as type information).
     """
-    definitions = evaluator.eval_element(stmt)
-    return chain.from_iterable(_execute_array_values(evaluator, d) for d in definitions)
+    definitions = module_context.eval_node(stmt)
+    return unite(_execute_array_values(module_context.evaluator, d) for d in definitions)
 
 
 def _execute_array_values(evaluator, array):
@@ -165,35 +175,33 @@ def _execute_array_values(evaluator, array):
     """
     if isinstance(array, ArrayLiteralContext):
         values = []
-        for types in array.py__iter__():
-            objects = set(chain.from_iterable(_execute_array_values(evaluator, typ) for typ in types))
-            values.append(AlreadyEvaluated(objects))
-        return [FakeSequence(evaluator, values, array.type)]
+        for lazy_context in array.py__iter__():
+            objects = unite(_execute_array_values(evaluator, typ) for typ in lazy_context.infer())
+            values.append(context.LazyKnownContexts(objects))
+        return set([FakeSequence(evaluator, array.array_type, values)])
     else:
-        return evaluator.execute(array)
+        return array.execute_evaluated()
 
 
-@memoize_default(None, evaluator_is_first_arg=True)
-def follow_param(evaluator, param):
+@memoize_default()
+def follow_param(module_context, param):
     def eval_docstring(docstring):
         return set(
             [p for param_str in _search_param_in_docstr(docstring, str(param.name))
-                for p in _evaluate_for_statement_string(evaluator, param_str, module)]
+                for p in _evaluate_for_statement_string(module_context, param_str)]
         )
     func = param.parent_function
-    module = param.get_parent_until()
-
     types = eval_docstring(func.raw_doc)
     if func.name.value == '__init__':
-        cls = func.get_parent_until(Class)
+        cls = search_ancestor(func, 'classdef')
         if cls.type == 'classdef':
             types |= eval_docstring(cls.raw_doc)
 
     return types
 
 
-@memoize_default(None, evaluator_is_first_arg=True)
-def find_return_types(evaluator, func):
+@memoize_default()
+def find_return_types(module_context, func):
     def search_return_in_docstr(code):
         for p in DOCSTRING_RETURN_PATTERNS:
             match = p.search(code)
@@ -201,4 +209,4 @@ def find_return_types(evaluator, func):
                 return _strip_rst_role(match.group(1))
 
     type_str = search_return_in_docstr(func.raw_doc)
-    return _evaluate_for_statement_string(evaluator, type_str, func.get_parent_until())
+    return _evaluate_for_statement_string(module_context, type_str)