Add codemod to convert typing.Union to |

libcst/codemod/commands/convert_union_to_or.py

@@ -0,0 +1,56 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# pyre-strict
+
+import libcst as cst
+from libcst.codemod import VisitorBasedCodemodCommand
+from libcst.codemod.visitors import RemoveImportsVisitor
+from libcst.metadata import QualifiedName, QualifiedNameProvider, QualifiedNameSource
+
+
+class ConvertUnionToOrCommand(VisitorBasedCodemodCommand):
+    DESCRIPTION: str = "Convert `Union[A, B]` to `A | B` in Python 3.10+"
+
+    METADATA_DEPENDENCIES = (QualifiedNameProvider,)
+
+    def leave_Subscript(
+        self, original_node: cst.Subscript, updated_node: cst.Subscript
+    ) -> cst.BaseExpression:
+        """
+        Given a subscript, check if it's a Union - if so, either flatten the members
+        into a nested BitOr (if multiple members) or unwrap the type (if only one member).
+        """
+        if QualifiedNameProvider.has_name(
+            self,
+            original_node,
+            QualifiedName(name="typing.Union", source=QualifiedNameSource.IMPORT),
+        ):
+            types = [
+                cst.ensure_type(
+                    cst.ensure_type(s, cst.SubscriptElement).slice, cst.Index
+                ).value
+                for s in updated_node.slice
+            ]
+            if len(types) == 1:
+                replacement = types[0]
+            else:
+                replacement = cst.BinaryOperation(
+                    left=types[0], right=types[1], operator=cst.BitOr()
+                )
+                for type_ in types[2:]:
+                    replacement = cst.BinaryOperation(
+                        left=replacement, right=type_, operator=cst.BitOr()
+                    )
+            return replacement
+        return updated_node
+
+    def leave_Module(
+        self, original_node: cst.Module, updated_node: cst.Module
+    ) -> cst.Module:
+        RemoveImportsVisitor.remove_unused_import(
+            self.context, module="typing", obj="Union"
+        )
+        return updated_node

libcst/codemod/commands/tests/test_convert_union_to_or.py

@@ -0,0 +1,86 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# pyre-strict
+
+from libcst.codemod import CodemodTest
+from libcst.codemod.commands.convert_union_to_or import ConvertUnionToOrCommand
+
+
+class TestConvertUnionToOrCommand(CodemodTest):
+    TRANSFORM = ConvertUnionToOrCommand
+
+    def test_simple_union(self) -> None:
+        before = """
+            from typing import Union
+            x: Union[int, str]
+        """
+        after = """
+            x: int | str
+        """
+        self.assertCodemod(before, after)
+
+    def test_nested_union(self) -> None:
+        before = """
+            from typing import Union
+            x: Union[int, Union[str, float]]
+        """
+        after = """
+            x: int | str | float
+        """
+        self.assertCodemod(before, after)
+
+    def test_single_type_union(self) -> None:
+        before = """
+            from typing import Union
+            x: Union[int]
+        """
+        after = """
+            x: int
+        """
+        self.assertCodemod(before, after)
+
+    def test_union_with_alias(self) -> None:
+        before = """
+            import typing as t
+            x: t.Union[int, str]
+        """
+        after = """
+            import typing as t
+            x: int | str
+        """
+        self.assertCodemod(before, after)
+
+    def test_union_with_unused_import(self) -> None:
+        before = """
+            from typing import Union, List
+            x: Union[int, str]
+        """
+        after = """
+            from typing import List
+            x: int | str
+        """
+        self.assertCodemod(before, after)
+
+    def test_union_no_import(self) -> None:
+        before = """
+            x: Union[int, str]
+        """
+        after = """
+            x: Union[int, str]
+        """
+        self.assertCodemod(before, after)
+
+    def test_union_in_function(self) -> None:
+        before = """
+            from typing import Union
+            def foo(x: Union[int, str]) -> Union[float, None]:
+                ...
+        """
+        after = """
+            def foo(x: int | str) -> float | None:
+                ...
+        """
+        self.assertCodemod(before, after)