[ExpressionContextProvider] Make subscript values always have a LOAD …

…context Fixes Instagram#318.
zsol · Jun 18, 2020 · fd938b3 · fd938b3
1 parent 5992a7d
commit fd938b3
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 5 deletions.
diff --git a/libcst/codemod/commands/tests/test_remove_unused_imports.py b/libcst/codemod/commands/tests/test_remove_unused_imports.py
@@ -76,3 +76,10 @@ def foo() -> None:
         """
 
         self.assertCodemod(before, after)
+
+    def test_access_in_assignment(self) -> None:
+        before = """
+            from a import b
+            b(0)[x] = False
+        """
+        self.assertCodemod(before, before)
diff --git a/libcst/metadata/expression_context_provider.py b/libcst/metadata/expression_context_provider.py
@@ -133,7 +133,9 @@ def visit_Attribute(self, node: cst.Attribute) -> bool:
 
     def visit_Subscript(self, node: cst.Subscript) -> bool:
         self.provider.set_metadata(node, self.context)
-        node.value.visit(self)
+        node.value.visit(
+            ExpressionContextVisitor(self.provider, ExpressionContext.LOAD)
+        )
         slice = node.slice
         if isinstance(slice, Sequence):
             for sli in slice:

diff --git a/libcst/metadata/tests/test_expression_context_provider.py b/libcst/metadata/tests/test_expression_context_provider.py
@@ -43,6 +43,7 @@ def visit_Name(self, node: cst.Name) -> None:
         self.test.assertEqual(
             self.get_metadata(ExpressionContextProvider, node, None),
             self.name_to_context[node.value],
+            f"Context doesn't match for Name {node.value}",
         )
 
     def visit_Attribute(self, node: cst.Attribute) -> None:
@@ -148,7 +149,7 @@ def test_assign_with_subscript(self) -> None:
             DependentVisitor(
                 test=self,
                 name_to_context={
-                    "a": ExpressionContext.STORE,
+                    "a": ExpressionContext.LOAD,
                     "b": ExpressionContext.LOAD,
                     "c": ExpressionContext.LOAD,
                     "d": ExpressionContext.LOAD,
@@ -226,7 +227,7 @@ def test_del_with_subscript(self) -> None:
             DependentVisitor(
                 test=self,
                 name_to_context={
-                    "a": ExpressionContext.DEL,
+                    "a": ExpressionContext.LOAD,
                     "b": ExpressionContext.LOAD,
                 },
                 subscript_to_context={"a[b]": ExpressionContext.DEL},
@@ -278,7 +279,7 @@ def test_nested_tuple_with_assign(self) -> None:
             )
         )
 
-    def test_list_with_assing(self) -> None:
+    def test_list_with_assign(self) -> None:
         wrapper = MetadataWrapper(parse_module("[a] = [b]"))
         wrapper.visit(
             DependentVisitor(
@@ -294,7 +295,7 @@ def test_list_with_assing(self) -> None:
             )
         )
 
-    def test_nested_list_with_assing(self) -> None:
+    def test_nested_list_with_assign(self) -> None:
         wrapper = MetadataWrapper(parse_module("[[a, b], c] = [[d, e], f]"))
         wrapper.visit(
             DependentVisitor(
@@ -316,6 +317,21 @@ def test_nested_list_with_assing(self) -> None:
             )
         )
 
+    def test_expressions_with_assign(self) -> None:
+        wrapper = MetadataWrapper(parse_module("f(a)[b] = c"))
+        wrapper.visit(
+            DependentVisitor(
+                test=self,
+                name_to_context={
+                    "a": ExpressionContext.LOAD,
+                    "b": ExpressionContext.LOAD,
+                    "c": ExpressionContext.LOAD,
+                    "f": ExpressionContext.LOAD,
+                },
+                subscript_to_context={"f(a)[b]": ExpressionContext.STORE},
+            )
+        )
+
     def test_invalid_type_for_context(self) -> None:
         wrapper = MetadataWrapper(parse_module("a()"))
         wrapper.visit(