diff --git a/jedi/evaluate/arguments.py b/jedi/evaluate/arguments.py index 2642f487..8f66a6be 100644 --- a/jedi/evaluate/arguments.py +++ b/jedi/evaluate/arguments.py @@ -233,8 +233,8 @@ class TreeArguments(AbstractArguments): named_args.append((c[0].value, LazyTreeContext(self.context, c[2]),)) else: # Generator comprehension. # Include the brackets with the parent. - comp = iterable.GeneratorComprehension( - self._evaluator, self.context, self.argument_node.parent) + comp = iterable.ArgumentGeneratorComprehension( + self._evaluator, self.context, el) yield None, LazyKnownContext(comp) else: yield None, LazyTreeContext(self.context, el) @@ -246,10 +246,10 @@ class TreeArguments(AbstractArguments): def _as_tree_tuple_objects(self): for star_count, argument in unpack_arglist(self.argument_node): + default = None if argument.type == 'argument': - argument, default = argument.children[::2] - else: - default = None + if len(argument.children) == 3: # Keyword argument. + argument, default = argument.children[::2] yield argument, default, star_count def iter_calling_names_with_star(self): diff --git a/jedi/evaluate/context/iterable.py b/jedi/evaluate/context/iterable.py index 6ae64125..4fb30db7 100644 --- a/jedi/evaluate/context/iterable.py +++ b/jedi/evaluate/context/iterable.py @@ -127,22 +127,17 @@ class ComprehensionMixin(object): self._defining_context = defining_context self._atom = atom - def _get_comprehension(self): - "return 'a for a in b'" - # The atom contains a testlist_comp - return self._atom.children[1] - def _get_comp_for(self): - "return CompFor('for a in b')" - return self._get_comprehension().children[1] + """return CompFor('for a in b')""" + return self._atom.children[1].children[1] - def _eval_node(self, index=0): + def _entry_node(self): """ The first part `x + 1` of the list comprehension: [x + 1 for x in foo] """ - return self._get_comprehension().children[index] + return self._atom.children[1].children[0] @evaluator_method_cache() def _get_comp_for_context(self, parent_context, comp_for): @@ -174,9 +169,9 @@ class ComprehensionMixin(object): for result in self._nested(comp_fors[1:], context_): yield result except IndexError: - iterated = context_.eval_node(self._eval_node()) + iterated = context_.eval_node(self._entry_node()) if self.array_type == 'dict': - yield iterated, context_.eval_node(self._eval_node(2)) + yield iterated, context_.eval_node(self._value_node()) else: yield iterated @@ -253,7 +248,10 @@ class DictComprehension(_DictMixin, ComprehensionMixin, Sequence): array_type = u'dict' def _get_comp_for(self): - return self._get_comprehension().children[3] + return self._atom.children[1].children[3] + + def _value_node(self): + return self._atom.children[1].children[2] def py__iter__(self, contextualized_node=None): for keys, values in self._iterate(): @@ -307,6 +305,16 @@ class GeneratorComprehension(ComprehensionMixin, GeneratorBase): pass +class ArgumentGeneratorComprehension(ComprehensionMixin, GeneratorBase): + def _get_comp_for(self): + # Not actually an atom. But need to correct this comprehension madness + # anyway. + return self._atom.children[1] + + def _entry_node(self): + return self._atom.children[0] + + class SequenceLiteralContext(Sequence): _TUPLE_LIKE = 'testlist_star_expr', 'testlist', 'subscriptlist' mapping = {'(': u'tuple', diff --git a/test/completion/comprehensions.py b/test/completion/comprehensions.py index b2da01bc..e66858e0 100644 --- a/test/completion/comprehensions.py +++ b/test/completion/comprehensions.py @@ -223,3 +223,33 @@ next(iter({a for a in range(10)})) [int(str(x.value) for x in list def reset_missing_bracket(): pass + + +# ----------------- +# function calls +# ----------------- + +def foo(arg): + return arg + + +x = foo(x for x in [1]) + +#? int() +next(x) +#? +x[0] + +# While it's illegal to have more than one argument, when a generator +# expression is involved, it's still a valid parse tree and Jedi should still +# work (and especially not raise Exceptions). It's debatable wheter inferring +# values for invalid statements is a good idea, but not failing is a must. + +#? int() +next(foo(x for x in [1], 1)) + +def bar(x, y): + return y + +#? str() +next(bar(x for x in [1], x for x in ['']))