diff --git a/jedi/inference/filters.py b/jedi/inference/filters.py index 375711f8..4418ceaf 100644 --- a/jedi/inference/filters.py +++ b/jedi/inference/filters.py @@ -61,7 +61,9 @@ def _get_definition_names(used_names, name_key): return for_module[name_key] except KeyError: names = used_names.get(name_key, ()) - result = for_module[name_key] = tuple(name for name in names if name.is_definition()) + result = for_module[name_key] = tuple( + name for name in names if name.is_definition() + ) return result diff --git a/jedi/inference/syntax_tree.py b/jedi/inference/syntax_tree.py index b5faecfd..86eb0d47 100644 --- a/jedi/inference/syntax_tree.py +++ b/jedi/inference/syntax_tree.py @@ -18,6 +18,7 @@ from jedi.inference import imports from jedi.inference import arguments from jedi.inference.value import ClassValue, FunctionValue from jedi.inference.value import iterable +from jedi.inference.value.dynamic_arrays import ListModification, DictModification from jedi.inference.value import TreeInstance from jedi.inference.helpers import is_string, is_literal, is_number from jedi.inference.compiled.access import COMPARISON_OPERATORS @@ -291,8 +292,24 @@ def _infer_expr_stmt(context, stmt, seek_name=None): names are defined in the statement, `seek_name` returns the result for this name. + expr_stmt: testlist_star_expr (annassign | augassign (yield_expr|testlist) | + ('=' (yield_expr|testlist_star_expr))*) + annassign: ':' test ['=' test] + augassign: ('+=' | '-=' | '*=' | '@=' | '/=' | '%=' | '&=' | '|=' | '^=' | + '<<=' | '>>=' | '**=' | '//=') + :param stmt: A `tree.ExprStmt`. """ + def check_setitem(stmt): + atom_expr = stmt.children[0] + if atom_expr.type not in ('atom_expr', 'power'): + return False, None + name = atom_expr.children[0] + if name.type != 'name' or len(atom_expr.children) != 2: + return False, None + trailer = atom_expr.children[-1] + return trailer.children[0] == '[', trailer.children[1] + debug.dbg('infer_expr_stmt %s (%s)', stmt, seek_name) rhs = stmt.get_rhs() value_set = context.infer_node(rhs) @@ -302,31 +319,44 @@ def _infer_expr_stmt(context, stmt, seek_name=None): value_set = check_tuple_assignments(c_node, value_set) first_operator = next(stmt.yield_operators(), None) - if first_operator not in ('=', None) and first_operator.type == 'operator': + is_setitem, subscriptlist = check_setitem(stmt) + is_annassign = first_operator not in ('=', None) and first_operator.type == 'operator' + if is_annassign or is_setitem: # `=` is always the last character in aug assignments -> -1 - operator = copy.copy(first_operator) - operator.value = operator.value[:-1] name = stmt.get_defined_names()[0].value - left = context.py__getattribute__(name, position=stmt.start_pos) + left_values = context.py__getattribute__(name, position=stmt.start_pos) - for_stmt = tree.search_ancestor(stmt, 'for_stmt') - if for_stmt is not None and for_stmt.type == 'for_stmt' and value_set \ - and parser_utils.for_stmt_defines_one_name(for_stmt): - # Iterate through result and add the values, that's possible - # only in for loops without clutter, because they are - # predictable. Also only do it, if the variable is not a tuple. - node = for_stmt.get_testlist() - cn = ContextualizedNode(context, node) - ordered = list(cn.infer().iterate(cn)) + if is_setitem: + def to_mod(v): + c = ContextualizedNode(context, subscriptlist) + if v.array_type == 'dict': + return DictModification(v, value_set, c) + elif v.array_type == 'list': + return ListModification(v, value_set, c) + return v - for lazy_value in ordered: - dct = {for_stmt.children[1].value: lazy_value.infer()} - with context.predefine_names(for_stmt, dct): - t = context.infer_node(rhs) - left = _infer_comparison(context, left, operator, t) - value_set = left + value_set = ValueSet(to_mod(v) for v in left_values) else: - value_set = _infer_comparison(context, left, operator, value_set) + operator = copy.copy(first_operator) + operator.value = operator.value[:-1] + for_stmt = tree.search_ancestor(stmt, 'for_stmt') + if for_stmt is not None and for_stmt.type == 'for_stmt' and value_set \ + and parser_utils.for_stmt_defines_one_name(for_stmt): + # Iterate through result and add the values, that's possible + # only in for loops without clutter, because they are + # predictable. Also only do it, if the variable is not a tuple. + node = for_stmt.get_testlist() + cn = ContextualizedNode(context, node) + ordered = list(cn.infer().iterate(cn)) + + for lazy_value in ordered: + dct = {for_stmt.children[1].value: lazy_value.infer()} + with context.predefine_names(for_stmt, dct): + t = context.infer_node(rhs) + left_values = _infer_comparison(context, left_values, operator, t) + value_set = left_values + else: + value_set = _infer_comparison(context, left_values, operator, value_set) debug.dbg('infer_expr_stmt result %s', value_set) return value_set diff --git a/jedi/inference/value/dynamic_arrays.py b/jedi/inference/value/dynamic_arrays.py index 229a5dfb..f259821a 100644 --- a/jedi/inference/value/dynamic_arrays.py +++ b/jedi/inference/value/dynamic_arrays.py @@ -22,10 +22,14 @@ It is important to note that: from jedi import debug from jedi import settings from jedi.inference import recursion -from jedi.inference.base_value import ValueSet, NO_VALUES, HelperValueMixin +from jedi.inference.base_value import ValueSet, NO_VALUES, HelperValueMixin, \ + ValueWrapper +from jedi.inference.lazy_value import LazyKnownValues from jedi.inference.helpers import infer_call_of_leaf from jedi.inference.cache import inference_state_method_cache +_sentinel = object() + def check_array_additions(context, sequence): """ Just a mapper function for the internal _internal_check_array_additions """ @@ -162,3 +166,36 @@ class _ArrayInstance(HelperValueMixin): def iterate(self, contextualized_node=None, is_async=False): return self.py__iter__(contextualized_node) + + +class _Modification(ValueWrapper): + def __init__(self, wrapped_value, assigned_values, contextualized_key): + super(_Modification, self).__init__(wrapped_value) + self._assigned_values = assigned_values + self._contextualized_key = contextualized_key + + def py__getitem__(self, *args, **kwargs): + return self._wrapped_value.py__getitem__(*args, **kwargs) | self._assigned_values + + def py__simple_getitem__(self, index): + actual = [ + v.get_safe_value(_sentinel) + for v in self._contextualized_key.infer() + ] + if index in actual: + return self._assigned_values + return self._wrapped_value.py__simple_getitem__(index) + + +class DictModification(_Modification): + def py__iter__(self): + for lazy_context in self._wrapped_value.py__iter__(): + yield lazy_context + yield self._contextualized_key + + +class ListModification(_Modification): + def py__iter__(self): + for lazy_context in self._wrapped_value.py__iter__(): + yield lazy_context + yield LazyKnownValues(self._assigned_values) diff --git a/test/completion/dynamic_arrays.py b/test/completion/dynamic_arrays.py index 0f03c8fd..a594ed3c 100644 --- a/test/completion/dynamic_arrays.py +++ b/test/completion/dynamic_arrays.py @@ -307,3 +307,46 @@ lst.append('') #? float() int() str() lst[0] +# ----------------- +# list setitem +# ----------------- + +some_lst = [int] +some_lst[3] = str +#? int +some_lst[0] +#? str +some_lst[3] +#? int str +some_lst[2] + +some_lst[0] = tuple +#? tuple +some_lst[0] +#? int str tuple +some_lst[1] + +# ----------------- +# set setitem (should not work) +# ----------------- + +some_set = {int} +some_set[3] = str +#? int +some_set[0] +#? int +some_set[3] + +# ----------------- +# dict setitem +# ----------------- + +some_dct = {'a': float, 1: int} +some_dct['x'] = list +some_dct['y'] = tuple +#? list +some_dct['x'] +#? int float list tuple +some_dct['unknown'] +#? float +some_dct['a']