Implement patchedast for TypeVar syntax in function definition

Author: lieryan

CHANGELOG.md

@@ -2,7 +2,8 @@
 
 - #850 Update and pin black version in pre-commit and Github Actions
 - #851 Bump supported python version to up to Python 3.14
-- #852 Implement patchedast for TypeAlias
+- #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)
@@ -829,6 +832,14 @@ 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])
+        if node.default_value:
+            children.extend(["=", node.default_value])
+        self._handle(node, children)
+
 
 class _Source:
     def __init__(self, source):

ropetest/refactor/patchedasttest.py

@@ -37,6 +37,28 @@ 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 test_operator_support_completeness(self):
         ast_ops = {
             n.__name__
@@ -1524,6 +1546,97 @@ def test_type_alias(self):
             "Subscript",
         ])
 
+    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",
+            ],
+        )
+
+    def test_type_var_with_constraint(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",
+        ])
+
+    def test_type_var_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(checker)
+
+        checker.check_children("TypeVar", [
+            "T",
+            "",
+            ":",
+            " ",
+            "Tuple",
+        ])
+
+    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",
+        ])
+
 
 class _ResultChecker:
     def __init__(self, test_case, ast):