Add TypeVarTuple node

Author: jacobtylerwalls

astroid/inference.py

@@ -95,6 +95,7 @@ def infer_end(
 nodes.TypeAlias._infer = infer_end # type: ignore[assignment]
 nodes.TypeVar._infer = infer_end # type: ignore[assignment]
 nodes.ParamSpec._infer = infer_end # type: ignore[assignment]
+nodes.TypeVarTuple._infer = infer_end # type: ignore[assignment]
 
 
 def _infer_sequence_helper(

astroid/nodes/__init__.py

@@ -86,6 +86,7 @@
  Tuple,
  TypeAlias,
  TypeVar,
+ TypeVarTuple,
  UnaryOp,
  Unknown,
  While,
@@ -184,6 +185,7 @@
  NodeNG,
  Nonlocal,
  ParamSpec,
+ TypeVarTuple,
  Pass,
  Pattern,
  Raise,
@@ -294,6 +296,7 @@
  "Tuple",
  "TypeAlias",
  "TypeVar",
+ "TypeVarTuple",
  "UnaryOp",
  "Unknown",
  "unpack_infer",

astroid/nodes/as_string.py

@@ -531,6 +531,10 @@ def visit_typevar(self, node: nodes.TypeVar) -> str:
  """return an astroid.TypeVar node as string"""
  return node.name.accept(self) if node.name else "_"
 
+ def visit_typevartuple(self, node: nodes.TypeVarTuple) -> str:
+ """return an astroid.TypeVarTuple node as string"""
+ return "*" + node.name.accept(self) if node.name else ""
+
  def visit_unaryop(self, node) -> str:
  """return an astroid.UnaryOp node as string"""
  if node.op == "not":

astroid/nodes/node_classes.py

@@ -3376,7 +3376,7 @@ def __init__(
  end_col_offset: int | None = None,
  ) -> None:
  self.name: AssignName | None
- self.type_params: list[TypeVar, ParamSpec]
+ self.type_params: list[TypeVar, ParamSpec, TypeVarTuple]
  self.value: NodeNG
  super().__init__(
  lineno=lineno,
@@ -3390,7 +3390,7 @@ def postinit(
  self,
  *,
  name: AssignName | None,
- type_params: list[TypeVar, ParamSpec],
+ type_params: list[TypeVar, ParamSpec, TypeVarTuple],
  value: NodeNG,
  ) -> None:
  self.name = name
@@ -3448,6 +3448,37 @@ def postinit(self, *, name: AssignName | None, bound: NodeNG | None) -> None:
  self.bound = bound
 
 
+class TypeVarTuple(_base_nodes.AssignTypeNode):
+ """Class representing a :class:`ast.TypeVarTuple` node.
+
+ >>> import astroid
+ >>> node = astroid.extract_node('type Alias[*Ts] = tuple[*Ts]')
+ >>> node.type_params[0]
+ <TypeVarTuple l.1 at 0x7f23b2e4e198>
+ """
+
+ 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: AssignName | 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: AssignName | None) -> None:
+ self.name = name
+
+
 class UnaryOp(NodeNG):
  """Class representing an :class:`ast.UnaryOp` node.
 

astroid/nodes/scoped_nodes/scoped_nodes.py

@@ -1130,7 +1130,7 @@ def __init__(
  self.body: list[NodeNG] = []
  """The contents of the function body."""
 
- self.type_params: list[nodes.TypeVar] = []
+ self.type_params: list[nodes.TypeVar, nodes.ParamSpec, nodes.TypeVarTuple] = []
  """PEP 695 (Python 3.12+) type params, e.g. first 'T' in def func[T]() -> T: ..."""
 
  self.instance_attrs: dict[str, list[NodeNG]] = {}
@@ -1846,7 +1846,7 @@ def __init__(
  self.is_dataclass: bool = False
  """Whether this class is a dataclass."""
 
- self.type_params: list[nodes.TypeVar] = []
+ self.type_params: list[nodes.TypeVar, nodes.ParamSpec, nodes.TypeVarTuple] = []
  """PEP 695 (Python 3.12+) type params, e.g. class MyClass[T]: ..."""
 
  super().__init__(

astroid/rebuilder.py

@@ -448,6 +448,12 @@ def visit(self, node: ast.TypeAlias, parent: NodeNG) -> nodes.TypeAlias:
  def visit(self, node: ast.TypeVar, parent: NodeNG) -> nodes.TypeVar:
  ...
 
+ @overload
+ def visit(
+ self, node: ast.TypeVarTuple, parent: NodeNG
+ ) -> nodes.TypeVarTuple:
+ ...
+
  @overload
  def visit(self, node: ast.UnaryOp, parent: NodeNG) -> nodes.UnaryOp:
  ...
@@ -1738,6 +1744,22 @@ def visit_typevar(self, node: ast.TypeVar, parent: NodeNG) -> nodes.TypeVar:
  )
  return newnode
 
+ def visit_typevartuple(
+ self, node: ast.TypeVarTuple, parent: NodeNG
+ ) -> nodes.TypeVarTuple:
+ """Visit a TypeVarTuple node by returning a fresh instance of it."""
+ newnode = nodes.TypeVarTuple(
+ lineno=node.lineno,
+ col_offset=node.col_offset,
+ end_lineno=node.end_lineno,
+ end_col_offset=node.end_col_offset,
+ parent=parent,
+ )
+ # Add AssignName node for 'node.name'
+ # https://bugs.python.org/issue43994
+ newnode.postinit(name=self.visit_assignname(node, newnode, node.name))
+ 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

@@ -82,6 +82,7 @@ Nodes
  astroid.nodes.Tuple
  astroid.nodes.TypeAlias
  astroid.nodes.TypeVar
+ astroid.nodes.TypeVarTuple
  astroid.nodes.UnaryOp
  astroid.nodes.Unknown
  astroid.nodes.While
@@ -235,6 +236,8 @@ Nodes
 
 .. autoclass:: astroid.nodes.TypeVar
 
+.. autoclass:: astroid.nodes.TypeVarTuple
+
 .. autoclass:: astroid.nodes.UnaryOp
 
 .. autoclass:: astroid.nodes.Unknown

tests/test_type_params.py

@@ -6,7 +6,7 @@
 
 from astroid import extract_node
 from astroid.const import PY312_PLUS
-from astroid.nodes import AssignName, ParamSpec, Subscript, TypeAlias, TypeVar
+from astroid.nodes import AssignName, ParamSpec, Subscript, TypeAlias, TypeVar, TypeVarTuple
 
 
 @pytest.mark.skipif(not PY312_PLUS, reason="Requires Python 3.12 or higher")
@@ -38,6 +38,17 @@ def test_type_param_spec() -> None:
  assert node.inferred()[0] is node
 
 
+@pytest.mark.skipif(not PY312_PLUS, reason="Requires Python 3.12 or higher")
+def test_type_var_tuple() -> None:
+ node = extract_node("type Alias[*Ts] = tuple[*Ts]")
+ params = node.type_params[0]
+ assert isinstance(params, TypeVarTuple)
+ assert isinstance(params.name, AssignName)
+ assert params.name.name == "Ts"
+
+ assert node.inferred()[0] is node
+
+
 @pytest.mark.skipif(not PY312_PLUS, reason="Requires Python 3.12 or higher")
 def test_type_param() -> None:
  func_node = extract_node("def func[T]() -> T: ...")