Add TypeAlias and TypeVar nodes (Python 3.12)

Author: jacobtylerwalls

astroid/nodes/__init__.py

@@ -83,6 +83,8 @@
     TryFinally,
     TryStar,
     Tuple,
+    TypeAlias,
+    TypeVar,
     UnaryOp,
     Unknown,
     While,
@@ -193,6 +195,8 @@
     TryFinally,
     TryStar,
     Tuple,
+    TypeAlias,
+    TypeVar,
     UnaryOp,
     Unknown,
     While,
@@ -285,6 +289,8 @@
     "TryFinally",
     "TryStar",
     "Tuple",
+    "TypeAlias",
+    "TypeVar",
     "UnaryOp",
     "Unknown",
     "unpack_infer",

astroid/nodes/as_string.py

@@ -178,6 +178,7 @@ def visit_classdef(self, node) -> str:
         args += [n.accept(self) for n in node.keywords]
         args_str = f"({', '.join(args)})" if args else ""
         docs = self._docs_dedent(node.doc_node)
+        # TODO: handle type_params
         return "\n\n{}class {}{}:{}\n{}\n".format(
             decorate, node.name, args_str, docs, self._stmt_list(node.body)
         )
@@ -330,6 +331,7 @@ def handle_functiondef(self, node, keyword) -> str:
         if node.returns:
             return_annotation = " -> " + node.returns.as_string()
             trailer = return_annotation + ":"
+        # TODO: handle type_params
         def_format = "\n%s%s %s(%s)%s%s\n%s"
         return def_format % (
             decorate,
@@ -517,6 +519,14 @@ def visit_tuple(self, node) -> str:
             return f"({node.elts[0].accept(self)}, )"
         return f"({', '.join(child.accept(self) for child in node.elts)})"
 
+    def visit_typealias(self, node: nodes.TypeAlias) -> str:
+        """return an astroid.TypeAlias node as string"""
+        return f"{node.value}{node.type_params or ''}"
+
+    def visit_typevar(self, node: nodes.TypeVar) -> str:
+        """return an astroid.TypeVar node as string"""
+        return node.name
+
     def visit_unaryop(self, node) -> str:
         """return an astroid.UnaryOp node as string"""
         if node.op == "not":

astroid/nodes/node_classes.py

@@ -20,7 +20,6 @@
     ClassVar,
     Literal,
     Optional,
-    TypeVar,
     Union,
 )
 
@@ -63,8 +62,8 @@ def _is_const(value) -> bool:
     return isinstance(value, tuple(CONST_CLS))
 
 
-_NodesT = TypeVar("_NodesT", bound=NodeNG)
-_BadOpMessageT = TypeVar("_BadOpMessageT", bound=util.BadOperationMessage)
+_NodesT = typing.TypeVar("_NodesT", bound=NodeNG)
+_BadOpMessageT = typing.TypeVar("_BadOpMessageT", bound=util.BadOperationMessage)
 
 AssignedStmtsPossibleNode = Union["List", "Tuple", "AssignName", "AssignAttr", None]
 AssignedStmtsCall = Callable[
@@ -3311,6 +3310,96 @@ def getitem(self, index, context: InferenceContext | None = None):
         return _container_getitem(self, self.elts, index, context=context)
 
 
+class TypeAlias(_base_nodes.AssignTypeNode):
+    """Class representing a :class:`ast.TypeAlias` node.
+
+    >>> import astroid
+    >>> node = astroid.extract_node('type Point = tuple[float, float]')
+    >>> node
+    <TypeAlias l.1 at 0x7f23b2e4e198>
+    """
+
+    _astroid_fields = ("type_params", "value")
+
+    def __init__(
+        self,
+        lineno: int | None = None,
+        col_offset: int | None = None,
+        parent: NodeNG | None = None,
+        *,
+        end_lineno: int | None = None,
+        end_col_offset: int | None = None,
+    ) -> None:
+        self.type_params: list[TypeVar]
+        self.value: NodeNG
+        super().__init__(
+            lineno=lineno,
+            col_offset=col_offset,
+            end_lineno=end_lineno,
+            end_col_offset=end_col_offset,
+            parent=parent,
+        )
+
+    def postinit(
+        self,
+        *,
+        type_params: list[TypeVar],
+        value: NodeNG,
+    ) -> None:
+        self.type_params = type_params
+        self.value = value
+
+    assigned_stmts: ClassVar[
+        Callable[
+            [
+                TypeAlias,
+                AssignName,
+                InferenceContext | None,
+                None,
+            ],
+            Generator[NodeNG, None, None],
+        ]
+    ]
+    """Returns the assigned statement (non inferred) according to the assignment type.
+    See astroid/protocols.py for actual implementation.
+    """
+
+
+class TypeVar(_base_nodes.AssignTypeNode):
+    """Class representing a :class:`ast.TypeVar` node.
+
+    >>> import astroid
+    >>> node = astroid.extract_node('type Point[T] = tuple[float, float]')
+    >>> node.type_params[0]
+    <TypeVar l.1 at 0x7f23b2e4e198>
+    """
+
+    _astroid_fields = ("bound",)
+
+    def __init__(
+        self,
+        lineno: int | None = None,
+        col_offset: int | None = None,
+        parent: NodeNG | None = None,
+        *,
+        end_lineno: int | None = None,
+        end_col_offset: int | None = None,
+    ) -> None:
+        self.name: str
+        self.bound: NodeNG | None
+        super().__init__(
+            lineno=lineno,
+            col_offset=col_offset,
+            end_lineno=end_lineno,
+            end_col_offset=end_col_offset,
+            parent=parent,
+        )
+
+    def postinit(self, *, name: str, bound: NodeNG | None) -> None:
+        self.name = name
+        self.bound = bound
+
+
 class UnaryOp(NodeNG):
     """Class representing an :class:`ast.UnaryOp` node.
 

astroid/nodes/scoped_nodes/scoped_nodes.py

@@ -1039,7 +1039,14 @@ class FunctionDef(
     <FunctionDef.my_func l.2 at 0x7f23b2e71e10>
     """
 
-    _astroid_fields = ("decorators", "args", "returns", "doc_node", "body")
+    _astroid_fields = (
+        "decorators",
+        "args",
+        "returns",
+        "type_params",
+        "doc_node",
+        "body",
+    )
     _multi_line_block_fields = ("body",)
     returns = None
 
@@ -1107,6 +1114,9 @@ def __init__(
         self.body: list[NodeNG] = []
         """The contents of the function body."""
 
+        self.type_params: list[nodes.TypeVar] = []
+        """PEP 695 (Python 3.12+) type params, e.g. first 'T' in def func[T]() -> T: ..."""
+
         self.instance_attrs: dict[str, list[NodeNG]] = {}
 
         super().__init__(
@@ -1131,6 +1141,7 @@ def postinit(
         *,
         position: Position | None = None,
         doc_node: Const | None = None,
+        type_params: list[nodes.TypeVar] | None = None,
     ):
         """Do some setup after initialisation.
 
@@ -1148,6 +1159,8 @@ def postinit(
             Position of function keyword(s) and name.
         :param doc_node:
             The doc node associated with this node.
+        :param type_params:
+            The type_params associated with this node.
         """
         self.args = args
         self.body = body
@@ -1157,6 +1170,7 @@ def postinit(
         self.type_comment_args = type_comment_args
         self.position = position
         self.doc_node = doc_node
+        self.type_params = type_params or []
 
     @cached_property
     def extra_decorators(self) -> list[node_classes.Call]:
@@ -1721,7 +1735,7 @@ def get_wrapping_class(node):
     return klass
 
 
-class ClassDef(
+class ClassDef(  # pylint: disable=too-many-instance-attributes
     _base_nodes.FilterStmtsBaseNode, LocalsDictNodeNG, _base_nodes.Statement
 ):
     """Class representing an :class:`ast.ClassDef` node.
@@ -1740,7 +1754,14 @@ def my_meth(self, arg):
     # by a raw factories
 
     # a dictionary of class instances attributes
-    _astroid_fields = ("decorators", "bases", "keywords", "doc_node", "body")  # name
+    _astroid_fields = (
+        "decorators",
+        "bases",
+        "keywords",
+        "doc_node",
+        "body",
+        "type_params",
+    )  # name
 
     decorators = None
     """The decorators that are applied to this class.
@@ -1807,6 +1828,9 @@ def __init__(
         self.is_dataclass: bool = False
         """Whether this class is a dataclass."""
 
+        self.type_params: list[nodes.TypeVar] = []
+        """PEP 695 (Python 3.12+) type params, e.g. class MyClass[T]: ..."""
+
         super().__init__(
             lineno=lineno,
             col_offset=col_offset,
@@ -1848,6 +1872,7 @@ def postinit(
         *,
         position: Position | None = None,
         doc_node: Const | None = None,
+        type_params: list[nodes.TypeVar] | None = None,
     ) -> None:
         if keywords is not None:
             self.keywords = keywords
@@ -1858,6 +1883,7 @@ def postinit(
         self._metaclass = metaclass
         self.position = position
         self.doc_node = doc_node
+        self.type_params = type_params or []
 
     def _newstyle_impl(self, context: InferenceContext | None = None):
         if context is None:

astroid/rebuilder.py

@@ -18,7 +18,7 @@
 
 from astroid import nodes
 from astroid._ast import ParserModule, get_parser_module, parse_function_type_comment
-from astroid.const import IS_PYPY, PY38, PY39_PLUS, Context
+from astroid.const import IS_PYPY, PY38, PY39_PLUS, PY312_PLUS, Context
 from astroid.manager import AstroidManager
 from astroid.nodes import NodeNG
 from astroid.nodes.utils import Position
@@ -432,6 +432,16 @@ def visit(self, node: ast.TryStar, parent: NodeNG) -> nodes.TryStar:
         def visit(self, node: ast.Tuple, parent: NodeNG) -> nodes.Tuple:
             ...
 
+        if sys.version_info >= (3, 12):
+
+            @overload
+            def visit(self, node: ast.TypeAlias, parent: NodeNG) -> nodes.TypeAlias:
+                ...
+
+            @overload
+            def visit(self, node: ast.TypeVar, parent: NodeNG) -> nodes.TypeVar:
+                ...
+
         @overload
         def visit(self, node: ast.UnaryOp, parent: NodeNG) -> nodes.UnaryOp:
             ...
@@ -870,6 +880,9 @@ def visit_classdef(
             ],
             position=self._get_position_info(node, newnode),
             doc_node=self.visit(doc_ast_node, newnode),
+            type_params=[self.visit(param, newnode) for param in node.type_params]
+            if PY312_PLUS
+            else [],
         )
         return newnode
 
@@ -1170,6 +1183,9 @@ def _visit_functiondef(
             type_comment_args=type_comment_args,
             position=self._get_position_info(node, newnode),
             doc_node=self.visit(doc_ast_node, newnode),
+            type_params=[self.visit(param, newnode) for param in node.type_params]
+            if PY312_PLUS
+            else [],
         )
         self._global_names.pop()
         return newnode
@@ -1669,6 +1685,33 @@ def visit_tuple(self, node: ast.Tuple, parent: NodeNG) -> nodes.Tuple:
         newnode.postinit([self.visit(child, newnode) for child in node.elts])
         return newnode
 
+    def visit_typealias(self, node: ast.TypeAlias, parent: NodeNG) -> nodes.TypeAlias:
+        """Visit a TypeAlias node by returning a fresh instance of it."""
+        newnode = nodes.TypeAlias(
+            lineno=node.lineno,
+            col_offset=node.col_offset,
+            end_lineno=node.end_lineno,
+            end_col_offset=node.end_col_offset,
+            parent=parent,
+        )
+        newnode.postinit(
+            type_params=[self.visit(p, newnode) for p in node.type_params],
+            value=self.visit(node.value, newnode),
+        )
+        return newnode
+
+    def visit_typevar(self, node: ast.TypeVar, parent: NodeNG) -> nodes.TypeVar:
+        """Visit a TypeVar node by returning a fresh instance of it."""
+        newnode = nodes.TypeVar(
+            lineno=node.lineno,
+            col_offset=node.col_offset,
+            end_lineno=node.end_lineno,
+            end_col_offset=node.end_col_offset,
+            parent=parent,
+        )
+        newnode.postinit(name=node.name, bound=self.visit(node.bound, newnode))
+        return newnode
+
     def visit_unaryop(self, node: ast.UnaryOp, parent: NodeNG) -> nodes.UnaryOp:
         """Visit a UnaryOp node by returning a fresh instance of it."""
         newnode = nodes.UnaryOp(

doc/api/astroid.nodes.rst

@@ -79,6 +79,8 @@ Nodes
    astroid.nodes.TryFinally
    astroid.nodes.TryStar
    astroid.nodes.Tuple
+   astroid.nodes.TypeAlias
+   astroid.nodes.TypeVar
    astroid.nodes.UnaryOp
    astroid.nodes.Unknown
    astroid.nodes.While
@@ -226,6 +228,10 @@ Nodes
 
 .. autoclass:: astroid.nodes.Tuple
 
+.. autoclass:: astroid.nodes.TypeAlias
+
+.. autoclass:: astroid.nodes.TypeVar
+
 .. autoclass:: astroid.nodes.UnaryOp
 
 .. autoclass:: astroid.nodes.Unknown