[Scope] Fix referencing of remaining objects in cast() (#422)

Instagram · Nov 30, 2020 · 2485d5a · 2485d5a
1 parent 02fc440
commit 2485d5a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 3 deletions.
diff --git a/libcst/metadata/scope_provider.py b/libcst/metadata/scope_provider.py
@@ -782,9 +782,12 @@ def visit_Call(self, node: cst.Call) -> Optional[bool]:
             return False
         if "typing.cast" in qnames:
             node.func.visit(self)
-            self.__in_type_hint.add(node)
             if len(node.args) > 0:
+                self.__in_type_hint.add(node)
                 node.args[0].visit(self)
+                self.__in_type_hint.discard(node)
+                for arg in node.args[1:]:
+                    arg.visit(self)
             return False
         return True
 

diff --git a/libcst/metadata/tests/test_scope_provider.py b/libcst/metadata/tests/test_scope_provider.py
@@ -1038,7 +1038,7 @@ def test_annotation_access(self) -> None:
         m, scopes = get_scope_metadata_provider(
             """
                 from typing import Literal, NewType, Optional, TypeVar, Callable, cast
-                from a import A, B, C, D, D2, E, E2, F, G, G2, H, I, J, K, K2
+                from a import A, B, C, D, D2, E, E2, F, G, G2, H, I, J, K, K2, L, M
                 def x(a: A):
                     pass
                 def y(b: "B"):
@@ -1054,7 +1054,8 @@ def z(c: Literal["C"]):
 
                 class Test(Generic[J]):
                     pass
-                casted = cast("K", "K2")
+                castedK = cast("K", "K2")
+                castedL = cast("L", M)
             """
         )
         imp = ensure_type(
@@ -1144,6 +1145,16 @@ class Test(Generic[J]):
         self.assertIsInstance(assignment, Assignment)
         self.assertEqual(len(assignment.references), 0)
 
+        assignment = list(scope["L"])[0]
+        self.assertIsInstance(assignment, Assignment)
+        self.assertEqual(len(assignment.references), 1)
+        references = list(assignment.references)
+
+        assignment = list(scope["M"])[0]
+        self.assertIsInstance(assignment, Assignment)
+        self.assertEqual(len(assignment.references), 1)
+        references = list(assignment.references)
+
     def test_node_of_scopes(self) -> None:
         m, scopes = get_scope_metadata_provider(
             """