Add a way how dict setitem can be understood

Needs the latest parso commits

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
 
 

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 is_setitem:
-        if for_stmt is not None and for_stmt.type == 'for_stmt' and value_set \
+            def to_mod(v):
-                and parser_utils.for_stmt_defines_one_name(for_stmt):
+                c = ContextualizedNode(context, subscriptlist)
-            # Iterate through result and add the values, that's possible
+                if v.array_type == 'dict':
-            # only in for loops without clutter, because they are
+                    return DictModification(v, value_set, c)
-            # predictable. Also only do it, if the variable is not a tuple.
+                elif v.array_type == 'list':
-            node = for_stmt.get_testlist()
+                    return ListModification(v, value_set, c)
-            cn = ContextualizedNode(context, node)
+                return v
-            ordered = list(cn.infer().iterate(cn))
 
-            for lazy_value in ordered:
+            value_set = ValueSet(to_mod(v) for v in left_values)
-                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
         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
 

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)

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']