Repository: Instagram/LibCST

Issue:
make @m.leave() recognize function return using new Union syntax | 
`@m.leave` throws an exception when function return use new union syntax:
```

    @m.call_if_inside(m.ImportFrom(module=m.Name(value="typing")))
    @m.leave(m.ImportFrom(module=m.Name(value="typing")))
    def test(
        self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
    ) -> cst.ImportFrom | cst.RemovalSentinel:
```

```
Traceback (most recent call last):
  File "/opt/instagram/virtualenv/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/opt/instagram/virtualenv/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/opt/instagram/virtualenv/lib/python3.7/site-packages/fbcode/libcst/tool.py", line 832, in <module>
    main(os.environ.get("LIBCST_TOOL_COMMAND_NAME", "libcst.tool"), sys.argv[1:])
  File "/opt/instagram/virtualenv/lib/python3.7/site-packages/fbcode/libcst/tool.py", line 827, in main
    return lookup.get(args.action or None, _invalid_command)(proc_name, command_args)
  File "/opt/instagram/virtualenv/lib/python3.7/site-packages/fbcode/libcst/tool.py", line 532, in _codemod_impl
    command_instance = command_class(CodemodContext(), **codemod_args)
  File "/opt/instagram/virtualenv/lib/python3.7/site-packages/fbcode/libcst/codemod/_visitor.py", line 29, in __init__
    MatcherDecoratableTransformer.__init__(self)
  File "/opt/instagram/virtualenv/lib/python3.7/site-packages/fbcode/libcst/matchers/_visitors.py", line 473, in __init__
    expected_none_return=False,
  File "/opt/instagram/virtualenv/lib/python3.7/site-packages/fbcode/libcst/matchers/_visitors.py", line 242, in _check_types
    expected_none=expected_none_return,
  File "/opt/instagram/virtualenv/lib/python3.7/site-packages/fbcode/libcst/matchers/_visitors.py", line 142, in _verify_return_annotation
    if issubclass(ret, annotation):
TypeError: issubclass() arg 1 must be a class
```

Old Union syntax works:
```
    @m.call_if_inside(m.ImportFrom(module=m.Name(value="typing")))
    @m.leave(m.ImportFrom(module=m.Name(value="typing")))
    def test(
        self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
    ) -> Union[cst.ImportFrom, cst.RemovalSentinel]:
```

Not using `@m.leave` works:
```
    @m.call_if_inside(m.ImportFrom(module=m.Name(value="typing")))
    def leave_ImportFrom(
        self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
    ) -> cst.ImportFrom | cst.RemovalSentinel:
```


Comments:

Patch:
diff --git a/libcst/matchers/_visitors.py b/libcst/matchers/_visitors.py
index 301e675..be50edf 100644
--- a/libcst/matchers/_visitors.py
+++ b/libcst/matchers/_visitors.py
@@ -79,8 +79,18 @@ def _get_possible_match_classes(matcher: BaseMatcherNode) -> List[Type[cst.CSTNo
         return [getattr(cst, matcher.__class__.__name__)]
 
 
-def _get_possible_annotated_classes(annotation: object) -> List[Type[object]]:
+def _annotation_looks_like_union(annotation: object) -> bool:
     if getattr(annotation, "__origin__", None) is Union:
+        return True
+    # support PEP-604 style unions introduced in Python 3.10
+    return (
+        annotation.__class__.__name__ == "Union"
+        and annotation.__class__.__module__ == "types"
+    )
+
+
+def _get_possible_annotated_classes(annotation: object) -> List[Type[object]]:
+    if _annotation_looks_like_union(annotation):
         return getattr(annotation, "__args__", [])
     else:
         return [cast(Type[object], annotation)]
diff --git a/libcst/metadata/scope_provider.py b/libcst/metadata/scope_provider.py
index 580cc11..b5c6ba6 100644
--- a/libcst/metadata/scope_provider.py
+++ b/libcst/metadata/scope_provider.py
@@ -45,7 +45,6 @@ _ASSIGNMENT_LIKE_NODES = (
     cst.AugAssign,
     cst.ClassDef,
     cst.CompFor,
-    cst.For,
     cst.FunctionDef,
     cst.Global,
     cst.Import,
@@ -125,7 +124,7 @@ class Access:
             for assignment in assignments
             if assignment.scope != self.scope or assignment._index < self.__index
         }
-        if not previous_assignments and assignments:
+        if not previous_assignments and assignments and self.scope.parent != self.scope:
             previous_assignments = self.scope.parent[name]
         self.__assignments |= previous_assignments
 
@@ -993,6 +992,14 @@ class ScopeVisitor(cst.CSTVisitor):
                 node.elt.visit(self)
         return False
 
+    def visit_For(self, node: cst.For) -> Optional[bool]:
+        node.target.visit(self)
+        self.scope._assignment_count += 1
+        for child in [node.iter, node.body, node.orelse, node.asynchronous]:
+            if child is not None:
+                child.visit(self)
+        return False
+
     def infer_accesses(self) -> None:
         # Aggregate access with the same name and batch add with set union as an optimization.
         # In worst case, all accesses (m) and assignments (n) refer to the same name,


Test patch:
diff --git a/libcst/matchers/tests/test_decorators.py b/libcst/matchers/tests/test_decorators.py
index c102f2a..b1ff3d0 100644
--- a/libcst/matchers/tests/test_decorators.py
+++ b/libcst/matchers/tests/test_decorators.py
@@ -6,6 +6,7 @@
 from ast import literal_eval
 from textwrap import dedent
 from typing import List, Set
+from unittest.mock import Mock
 
 import libcst as cst
 import libcst.matchers as m
@@ -993,3 +994,25 @@ class MatchersVisitLeaveDecoratorsTest(UnitTest):
 
         # We should have only visited a select number of nodes.
         self.assertEqual(visitor.visits, ['"baz"'])
+
+
+# This is meant to simulate `cst.ImportFrom | cst.RemovalSentinel` in py3.10
+FakeUnionClass: Mock = Mock()
+setattr(FakeUnionClass, "__name__", "Union")
+setattr(FakeUnionClass, "__module__", "types")
+FakeUnion: Mock = Mock()
+FakeUnion.__class__ = FakeUnionClass
+FakeUnion.__args__ = [cst.ImportFrom, cst.RemovalSentinel]
+
+
+class MatchersUnionDecoratorsTest(UnitTest):
+    def test_init_with_new_union_annotation(self) -> None:
+        class TransformerWithUnionReturnAnnotation(m.MatcherDecoratableTransformer):
+            @m.leave(m.ImportFrom(module=m.Name(value="typing")))
+            def test(
+                self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
+            ) -> FakeUnion:
+                pass
+
+        # assert that init (specifically _check_types on return annotation) passes
+        TransformerWithUnionReturnAnnotation()
diff --git a/libcst/metadata/tests/test_scope_provider.py b/libcst/metadata/tests/test_scope_provider.py
index c8a2b74..d1566aa 100644
--- a/libcst/metadata/tests/test_scope_provider.py
+++ b/libcst/metadata/tests/test_scope_provider.py
@@ -1568,6 +1568,47 @@ class ScopeProviderTest(UnitTest):
         self.assertEqual(len(a_comp_assignment.references), 1)
         self.assertEqual(list(a_comp_assignment.references)[0].node, comp.elt)
 
+    def test_for_scope_ordering(self) -> None:
+        m, scopes = get_scope_metadata_provider(
+            """
+            def f():
+                for x in []:
+                    x
+            class X:
+                def f():
+                    for x in []:
+                        x
+            """
+        )
+        for scope in scopes.values():
+            for acc in scope.accesses:
+                self.assertEqual(
+                    len(acc.referents),
+                    1,
+                    msg=(
+                        "Access for node has incorrect number of referents: "
+                        + f"{acc.node}"
+                    ),
+                )
+
+    def test_no_out_of_order_references_in_global_scope(self) -> None:
+        m, scopes = get_scope_metadata_provider(
+            """
+            x = y
+            y = 1
+            """
+        )
+        for scope in scopes.values():
+            for acc in scope.accesses:
+                self.assertEqual(
+                    len(acc.referents),
+                    0,
+                    msg=(
+                        "Access for node has incorrect number of referents: "
+                        + f"{acc.node}"
+                    ),
+                )
+
     def test_cast(self) -> None:
         with self.assertRaises(cst.ParserSyntaxError):
             m, scopes = get_scope_metadata_provider(