make @m.leave() recognize function return using new Union syntax |

[luciawlli]

@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:

[benhgreen]

Looks like the new Unions are not being identified because they don't define __origin__ as checked here. Is there a cleaner fix here than:

• inspecting __module__/__name__ on annotation.__class__ for types/typing.Union
• attempting to import types.Union (would fail on vanilla Python before 3.10) and checking annotation.__class__
• letting this check pass as long as __args__ is a list of classes

cc: @MaggieMoss

[jimmylai]

@benhgreen what do you get for annotation argument value with the new change?

[benhgreen]

@jimmylai In this example, the annotation would be cst.ImportFrom | cst.RemovalSentinel, which resolves to a types.Union.

[luciawlli]

We probably should bring this up to Maggie on her original post on Worksplace, to see if it would be beneficial to define __origin__ for the new Unions and see their response.

On LibCST side, we can probably go with inspecting __module__/__name__ on annotation.__class__ for types/typing.Union, which sounds like the best among the options (thanks for the thoughtful comment). In this case, is an import of types.Union also needed for the check? If so, probably #3 would be the simplest solution that gets the job done.

[benhgreen]

In this case, is an import of types.Union also needed for the check?

Since they're string literals, no import would be needed.