Merge pull request #854 from python-rope/lieryan-patchedast-type-alias-type-var

Implement patchedast for TypeAlias and TypeVar/TypeVar/ParamSpec syntax

Author: lieryan

CHANGELOG.md

@@ -1,7 +1,9 @@
 # **Upcoming release**
 
 - #850 Update and pin black version in pre-commit and Github Actions
-- #850 Bump supported python version to up to Python 3.14
+- #851 Bump supported python version to up to Python 3.14
+- #852 Implement patchedast handlers for TypeAlias 
+- #853 Implement patchedast handlers TypeVar
 
 # Release 1.14.0
 

rope/refactor/patchedast.py

@@ -491,6 +491,9 @@ def _handle_function_def_node(self, node, is_async):
             children.extend(("@", decorator))
         children.extend(["async", "def"] if is_async else ["def"])
         children.append(node.name)
+        type_params = getattr(node, "type_params", [])
+        if type_params:
+            children.extend(["[", *self._child_nodes(type_params, ","), "]"])
         children.extend(["(", node.args, ")"])
         children.append(":")
         children.extend(node.body)
@@ -825,6 +828,32 @@ def _MatchMapping(self, node):
         children.append("}")
         self._handle(node, children)
 
+    def _TypeAlias(self, node):
+        children = ["type", node.name, node.value]
+        self._handle(node, children)
+
+    def _TypeVar(self, node):
+        children = [node.name]
+        if node.bound:
+            children.extend([":", node.bound])
+        self._handle_default_value(node, children)
+        self._handle(node, children)
+
+    def _TypeVarTuple(self, node):
+        children = ["*", node.name]
+        self._handle_default_value(node, children)
+        self._handle(node, children)
+
+    def _ParamSpec(self, node):
+        children = ["**", node.name]
+        self._handle_default_value(node, children)
+        self._handle(node, children)
+
+    def _handle_default_value(self, node, children):
+        default_value = getattr(node, "default_value", None)
+        if default_value:
+            children.extend(["=", default_value])
+
 
 class _Source:
     def __init__(self, source):

ropetest/refactor/patchedasttest.py

@@ -37,6 +37,60 @@ def assert_single_case_match_block(self, checker, match_type):
             "Expr",
         ])
 
+    def assert_function_def_has_one_type_var(self, checker):
+        checker.check_children(
+            "FunctionDef",
+            [
+                "def",
+                " ",
+                "foo",
+                "",
+                "[", "", "TypeVar", "", "]",
+                "",
+                "(", "", "arguments", "", ")",
+                "",
+                ":",
+                "\n    ",
+                "Pass",
+            ],
+        )
+
+    def assert_function_def_has_one_type_var_tuple(self, checker):
+        checker.check_children(
+            "FunctionDef",
+            [
+                "def",
+                " ",
+                "foo",
+                "",
+                "[", "", "TypeVarTuple", "", "]",
+                "",
+                "(", "", "arguments", "", ")",
+                "",
+                ":",
+                "\n    ",
+                "Pass",
+            ],
+        )
+
+    def assert_function_def_has_one_param_spec(self, checker):
+        checker.check_children(
+            "FunctionDef",
+            [
+                "def",
+                " ",
+                "foo",
+                "",
+                "[", "", "ParamSpec", "", "]",
+                "",
+                "(", "", "arguments", "", ")",
+                "",
+                ":",
+                "\n    ",
+                "Pass",
+            ],
+        )
+
     def test_operator_support_completeness(self):
         ast_ops = {
             n.__name__
@@ -1510,6 +1564,192 @@ def test_match_node_with_match_mapping_match_as(self):
             "}",
         ])
 
+    @testutils.only_for_versions_higher("3.12")
+    def test_type_alias(self):
+        source = dedent("""\
+            type Point = tuple[float, float]
+        """)
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+        checker.check_children("TypeAlias", [
+            "type",
+            " ",
+            "Name",
+            " = ",
+            "Subscript",
+        ])
+
+    @testutils.only_for_versions_higher("3.12")
+    def test_type_var_simple(self):
+        source = dedent("""\
+            def foo[S, T](x):
+                pass
+        """)
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+
+        checker.check_children(
+            "FunctionDef",
+            [
+                "def",
+                " ",
+                "foo",
+                "",
+                "[",
+                "",
+                "TypeVar",
+                "",
+                ",",
+                " ",
+                "TypeVar",
+                "",
+                "]",
+                "",
+                "(", "", "arguments", "", ")",
+                "",
+                ":",
+                "\n    ",
+                "Pass",
+            ],
+        )
+
+    @testutils.only_for_versions_higher("3.13")
+    def test_type_var_with_default_value(self):
+        source = dedent("""\
+            def foo[T = D](x):
+                pass
+        """)
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+
+        self.assert_function_def_has_one_type_var(checker)
+
+        checker.check_children("TypeVar", [
+            "T",
+            " ",
+            "=",
+            " ",
+            "Name",
+        ])
+
+    @testutils.only_for_versions_higher("3.12")
+    def test_type_var_with_constraint(self):
+        source = dedent("""\
+            def foo[T: (A, B)](x):
+                pass
+        """)
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+
+        self.assert_function_def_has_one_type_var(checker)
+
+        checker.check_children("TypeVar", [
+            "T",
+            "",
+            ":",
+            " ",
+            "Tuple",
+        ])
+
+    @testutils.only_for_versions_higher("3.13")
+    def test_type_var_with_constraint_and_default_value(self):
+        source = dedent("""\
+            def foo[T: (A, B) = D](x):
+                pass
+        """)
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+
+        self.assert_function_def_has_one_type_var(checker)
+
+        checker.check_children("TypeVar", [
+            "T",
+            "",
+            ":",
+            " ",
+            "Tuple",
+            " ",
+            "=",
+            " ",
+            "Name",
+        ])
+
+    @testutils.only_for_versions_higher("3.12")
+    def test_type_var_tuple_simple(self):
+        source = dedent("""\
+            def foo[*T](x):
+                pass
+        """)
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+
+        self.assert_function_def_has_one_type_var_tuple(checker)
+
+        checker.check_children("TypeVarTuple", [
+            "*",
+            "",
+            "T",
+        ])
+
+    @testutils.only_for_versions_higher("3.13")
+    def test_type_var_tuple_with_default_value(self):
+        source = dedent("""\
+            def foo[*T = (A, B)](x):
+                pass
+        """)
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+
+        self.assert_function_def_has_one_type_var_tuple(checker)
+
+        checker.check_children("TypeVarTuple", [
+            "*",
+            "",
+            "T",
+            " ",
+            "=",
+            " ",
+            "Tuple",
+        ])
+
+    @testutils.only_for_versions_higher("3.12")
+    def test_param_spec_simple(self):
+        source = dedent("""\
+            def foo[**T](x):
+                pass
+        """)
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+
+        self.assert_function_def_has_one_param_spec(checker)
+
+        checker.check_children("ParamSpec", [
+            "**",
+            "",
+            "T",
+        ])
+
+    @testutils.only_for_versions_higher("3.13")
+    def test_param_spec_with_default_value(self):
+        source = dedent("""\
+            def foo[**T = (A, B)](x):
+                pass
+        """)
+        ast_frag = patchedast.get_patched_ast(source, True)
+        checker = _ResultChecker(self, ast_frag)
+
+        self.assert_function_def_has_one_param_spec(checker)
+
+        checker.check_children("ParamSpec", [
+            "**",
+            "",
+            "T",
+            " ",
+            "=",
+            " ",
+            "Tuple",
+        ])
+
 
 class _ResultChecker:
     def __init__(self, test_case, ast):