🐛 fix: fix pep695 union import

Author: zrr1999

pyfuture/codemod/pep622/match.py

@@ -252,8 +252,21 @@ def test5():
         test_value = 123
         if test_value == y:
             print("other")
+
+    >>> module = cst.parse_module(\"""
+    ... def test6():
+    ...    for i in range(2):
+    ...        match i:
+    ...            case 0:
+    ...                yield 0
+    ...            case 1:
+    ...                yield 1
+    ... \""")
+    >>> new_module = transformer.transform_module(module)
     """
 
+    # TODO(zrr1999): Need to support nested
+
     METADATA_DEPENDENCIES = (ScopeProvider,)
 
     def __init__(self, context: CodemodContext) -> None:

pyfuture/codemod/pep695/type_parameters.py

@@ -52,13 +52,14 @@ def test(x: __test_T) -> __test_T:
     ... \""")
     >>> new_module = transformer.transform_module(module)
     >>> print(new_module.code)
-    from typing import TypeVar
+    from typing import TypeVar, Union
     def __wrapper_func_test():
         __test_T = TypeVar("__test_T", bound = Union[int, str])
         def test(x: __test_T) -> __test_T:
             return x
         return test
     test = __wrapper_func_test()
+    >>> transformer = TransformTypeParametersCommand(CodemodContext())
     >>> module = cst.parse_module(\"""
     ... class Test[T: int]:
     ...     def test[P: str](self, x: T, y: P) -> tuple[T, P]:
@@ -94,8 +95,7 @@ def remove_type_parameters[T: FunctionDef | ClassDef](
             new_name = type_param.param.name.with_changes(value=f"{prefix}{type_param.param.name.value}{suffix}")
 
             AddImportsVisitor.add_needed_import(self.context, "typing", type_param.param.__class__.__name__)
-            AddImportsVisitor.add_needed_import(self.context, "typing", "Union")
-            statements.append(gen_type_param(type_param.param, new_name))
+            statements.append(gen_type_param(type_param.param, new_name, self.context))
             slices.append(
                 SubscriptElement(
                     slice=Index(value=new_name),

pyfuture/codemod/utils.py

@@ -4,7 +4,8 @@
 from enum import Enum
 
 import libcst as cst
-from libcst.codemod import Codemod
+from libcst.codemod import Codemod, CodemodContext
+from libcst.codemod.visitors import AddImportsVisitor
 
 
 class RuleSet(Enum):
@@ -102,7 +103,9 @@ def _split_bit_or(_op: cst.BinaryOperation) -> list[cst.BaseExpression]:
 
 
 def gen_type_param(
-    type_param: cst.TypeVar | cst.TypeVarTuple | cst.ParamSpec, type_name: cst.Name | None = None
+    type_param: cst.TypeVar | cst.TypeVarTuple | cst.ParamSpec,
+    type_name: cst.Name | None = None,
+    context: CodemodContext | None = None,
 ) -> cst.SimpleStatementLine:
     """
     To generate type parameter definition statement.
@@ -128,6 +131,12 @@ def gen_type_param(
             if bound is not None:
                 if isinstance(bound, cst.BinaryOperation):
                     bound = transform_bit_or(bound) or bound
+                    if (
+                        isinstance(bound, cst.Subscript)
+                        and isinstance(bound_value := bound.value, cst.Name)
+                        and context is not None
+                    ):
+                        AddImportsVisitor.add_needed_import(context, "typing", bound_value.value)
                 args.append(cst.Arg(bound, keyword=cst.Name("bound")))
             return cst.SimpleStatementLine(
                 [