✨ feat: support PEP604 (#11)

* ♻️ refactor: use new structure

* ♻ refactor:  use pep group

* ✅ test:  add more pep695 case

* 🔧 config: use pytest-cov

* 🐛 fix: shield space in libcst

* ✨ feat: support pep604

* 🐛 fix: fix deps in workflow

* 🐛 fix: use m.matches instead of cst.ensure_type

.github/workflows/codestyle.yml

@@ -19,9 +19,10 @@ jobs:
         with:
           python-version: "3.12"
 
-      - name: Install pyfuture
+      - name: Install dependencies
         run: |
           pip install -e .[all]
+          pip install pytest
 
       - name: Run pre-commit
         uses: pre-commit/action@v3.0.0

.github/workflows/coverage.yml

@@ -27,8 +27,7 @@ jobs:
 
       - name: Run tests and coverage
         run: |
-          pdm run coverage run -m pytest tests
-          pdm run coverage html --show-contexts --title "Coverage for ${{ github.sha }}"
+          pdm run cov --title "Coverage for ${{ github.sha }}"
 
       - name: Store coverage HTML
         uses: actions/upload-artifact@v3

pyfuture/__init__.py

@@ -1 +1,2 @@
 from .__version__ import __version__
+from .utils import apply_transformer, transfer_code, transfer_file

pyfuture/codemod/__init__.py

@@ -1,2 +1,2 @@
-from .transform_match import TransformMatchCommand
-from .transform_type_parameters import TransformTypeParametersCommand
+from .pep622 import TransformMatchCommand
+from .pep695 import TransformTypeParametersCommand

pyfuture/codemod/pep604/__init__.py

@@ -0,0 +1 @@
+from .union_types import TransformUnionTypesCommand

pyfuture/codemod/pep604/union_types.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import libcst as cst
+from libcst import matchers as m
+from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
+from libcst.codemod.visitors import AddImportsVisitor
+from libcst.metadata import ScopeProvider
+
+
+class TransformUnionTypesCommand(VisitorBasedCodemodCommand):
+    METADATA_DEPENDENCIES = (ScopeProvider,)
+
+    def __init__(self, context: CodemodContext) -> None:
+        super().__init__(context)
+
+    def transform_union(self, op: cst.BinaryOperation) -> cst.Subscript | None:
+        if not isinstance(op.operator, cst.BitOr):
+            return None
+        if isinstance((left := op.left), cst.BinaryOperation):
+            left = self.transform_union(left) or left
+        if isinstance((right := op.right), cst.BinaryOperation):
+            right = self.transform_union(right) or right
+        slices = [
+            cst.SubscriptElement(
+                slice=cst.Index(value=left),
+            ),
+            cst.SubscriptElement(
+                slice=cst.Index(value=right),
+            ),
+        ]
+        return cst.Subscript(
+            value=cst.Name(
+                value="Union",
+                lpar=[],
+                rpar=[],
+            ),
+            slice=slices,
+        )
+
+    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call):
+        if not m.matches(original_node.func, m.Name("isinstance") | m.Name("issubclass")):
+            return updated_node
+        args = original_node.args
+
+        if (
+            isinstance(cls_info := args[1].value, cst.BinaryOperation)
+            and (cls_info := self.transform_union(cls_info)) is not None
+        ):
+            return updated_node.with_changes(
+                args=[
+                    args[0],
+                    cst.Arg(cls_info),
+                ]
+            )
+
+        return updated_node
+
+    def leave_Annotation(self, original_node: cst.Annotation, updated_node: cst.Annotation):
+        if (
+            isinstance((op := original_node.annotation), cst.BinaryOperation)
+            and (new_annotation := self.transform_union(op)) is not None
+        ):
+            AddImportsVisitor.add_needed_import(self.context, "typing", "Union")
+            return updated_node.with_changes(annotation=new_annotation)
+        return updated_node

pyfuture/codemod/pep622/__init__.py

@@ -0,0 +1 @@
+from .match import TransformMatchCommand

pyfuture/codemod/transform_match.py → pyfuture/codemod/pep622/match.py

@@ -16,7 +16,7 @@
 )
 from libcst.metadata import FunctionScope, ScopeProvider
 
-from ..transformer import ReplaceTransformer
+from ...transformer import ReplaceTransformer
 
 
 def match_selector(left: cst.BaseExpression, case: cst.MatchCase):

pyfuture/codemod/pep695/__init__.py

@@ -0,0 +1 @@
+from .type_parameters import TransformTypeParametersCommand

pyfuture/codemod/transform_type_parameters.py → pyfuture/codemod/pep695/type_parameters.py

@@ -23,8 +23,8 @@
 from libcst.codemod.visitors import AddImportsVisitor
 from libcst.metadata import Scope, ScopeProvider
 
-from ..transformer import ReplaceTransformer
-from .utils import gen_func_wrapper, gen_type_param
+from ...transformer import ReplaceTransformer
+from ..utils import gen_func_wrapper, gen_type_param
 
 
 class TransformTypeParametersCommand(VisitorBasedCodemodCommand):

pyfuture/codemod/utils.py

@@ -1,6 +1,41 @@
 from __future__ import annotations
 
+from enum import Enum
+from typing import Iterable
+
 import libcst as cst
+from libcst.codemod import Codemod
+
+
+class RuleSet(Enum):
+    # python 3.10+
+    pep604 = "pep604"
+    pep622 = "pep622"
+    # python 3.12+
+    pep695 = "pep695"
+
+
+def get_transformers(rule_sets: list[RuleSet] | RuleSet) -> Iterable[type[Codemod]]:
+    """
+    Get codemod transformers for specified rule set.
+    """
+    from .pep604 import TransformUnionTypesCommand
+    from .pep622 import TransformMatchCommand
+    from .pep695 import TransformTypeParametersCommand
+
+    if not isinstance(rule_sets, list):
+        rule_sets = [rule_sets]
+
+    for rule_set in rule_sets:
+        match rule_set:
+            case RuleSet.pep604:
+                yield TransformUnionTypesCommand
+            case RuleSet.pep622:
+                yield TransformMatchCommand
+            case RuleSet.pep695:
+                yield TransformTypeParametersCommand
+            case _:
+                raise ValueError(f"Unknown rule set: {rule_set}")
 
 
 def gen_type_param(

pyfuture/utils.py

@@ -1,21 +1,29 @@
 from __future__ import annotations
 
+import contextlib
+import io
 from pathlib import Path
 
 import libcst as cst
 from libcst.codemod import Codemod, CodemodContext
 
+from .codemod.utils import RuleSet, get_transformers
 
-def transform_code(
-    transformers: list[Codemod],
+
+def apply_transformer(
+    transformers: list[type[Codemod]],
     code: str,
 ) -> str:
     """
     Transform code with some transformers, and return the transformed code.
     """
-    module = cst.parse_module(code)
-    for transformer in transformers:
-        module = transformer.transform_module(module)
+    with contextlib.redirect_stdout(io.StringIO()):
+        module = cst.parse_module(code)
+        # while True:
+        for transformer in transformers:
+            module = transformer(CodemodContext()).transform_module(module)
+            # if code == module.code:
+            #     break
     return module.code
 
 
@@ -27,24 +35,14 @@ def transfer_code(
     """
     Transfer code to specified target version of python.
     """
-    from .codemod import TransformMatchCommand, TransformTypeParametersCommand
 
     assert target[0] == 3, "Only support python3"
     transformers = []
     if target[1] < 12:
-        transformers.extend(
-            [
-                TransformTypeParametersCommand(CodemodContext()),
-            ]
-        )
+        transformers.extend(get_transformers(RuleSet.pep695))
     if target[1] < 10:
-        transformers.extend(
-            [
-                TransformMatchCommand(CodemodContext()),
-            ]
-        )
-    # TODO: Add more codemods here
-    new_code = transform_code(
+        transformers.extend(get_transformers([RuleSet.pep622, RuleSet.pep604]))
+    new_code = apply_transformer(
         transformers=transformers,
         code=code,
     )

pyproject.toml

@@ -53,15 +53,17 @@ required-imports = ["from __future__ import annotations"]
 [tool.pyright]
 include = ["pyfuture"]
 exclude = [
-    "pdm_build.py",
-    # "pyfuture/codemod/transform_match.py"
+    "pdm_build.py"
 ]
 
 [tool.pdm.dev-dependencies]
 test = [
     "pytest",
+    "pytest-cov",
     "coverage"
 ]
+[tool.pdm.scripts]
+cov ={ composite = ["pytest tests --cov=pyfuture --cov-context=test", "coverage html --show-contexts {args}"]}
 
 [tool.pdm.version]
 source = "scm"

tests/codemod/test_pep695.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import libcst as cst
+import pytest
+from libcst.codemod import CodemodContext
+
+from pyfuture.codemod import TransformTypeParametersCommand
+
+
+def print_string(s: str):
+    cc = s.split("\n")
+    for c in cc[:-1]:
+        print(f"\"{c.replace("\"", "\\\"")}\\n\"")
+    print(f"\"{cc[-1].replace("\"", "\\\"")}\"")
+
+
+@pytest.mark.parametrize(
+    ("src", "expected"),
+    (
+        pytest.param(
+            "def test(x: int) -> int:\n" "    return x",
+            "def test(x: int) -> int:\n" "    return x",
+            id="no generic function",
+        ),
+        pytest.param(
+            "def test[T: int](x: T) -> T:\n" "    return x",
+            "from typing import TypeVar\n"
+            "\n"
+            "def __wrapper_func_test():\n"
+            '    __test_T = TypeVar("__test_T", bound = int)\n'
+            "    def test(x: __test_T) -> __test_T:\n"
+            "        return x\n"
+            "    return test\n"
+            "test = __wrapper_func_test()",
+            id="single bound function",
+        ),
+        pytest.param(
+            "def test[T: int | str](x: T) -> T:\n" "    return x",
+            "from typing import TypeVar\n\n"
+            "def __wrapper_func_test():\n"
+            '    __test_T = TypeVar("__test_T", bound = int | str)\n'
+            "    def test(x: __test_T) -> __test_T:\n"
+            "        return x\n"
+            "    return test\n"
+            "test = __wrapper_func_test()",
+            id="union bonud function",
+        ),
+    ),
+)
+def test_type_parameters(src: str, expected: str):
+    module = cst.parse_module(src)
+    new_module = TransformTypeParametersCommand(CodemodContext()).transform_module(module)
+    print_string(new_module.code)
+    assert new_module.code == expected