HEL/midas
1
0
Fork 0
Code Issues Pull Requests Actions Releases Activity

docs: add docstrings to generator submodule

Browse Source
This commit is contained in:
Louis Heredero
2026-07-07 12:09:44 +02:00
parent c6b8c0a145
commit ed07b01563
4 changed files with 499 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
midas/generator/collector.py
View File

@@ -5,20 +5,45 @@ from typing import Callable
import midas.ast.python as p
AssertionBuilder = Callable[..., ast.expr]
"""A callback function which builds an assertion test given some input expressions"""
@dataclass
class Assertion:
"""Runtime assertion to generate, bound to an expression"""
bound_expr: p.Expr
"""The expression the assertion is bound to"""
inputs: list[p.Expr]
"""
Expressions needed for the assertion
Each expression will be converted by the generator and passed as individual
arguments to `builder`
"""
builder: AssertionBuilder
"""The callback to build the assertion test given converted expression from `inputs`"""
message: str
"""The assertion message"""
def is_bound_to(self, expr: p.Expr) -> bool:
"""Check whether this assertion is bound to the given expression
Args:
expr (p.Expr): the expression
Returns:
bool: whether this assertion is bound to `expr`
"""
return expr == self.bound_expr
class AssertionCollector:
"""Helper class to collect assertions from outside the generator"""
def __init__(self):
self.assertions: list[Assertion] = []
self.definitions: dict[str, ast.stmt] = {}
@@ -30,6 +55,15 @@ class AssertionCollector:
builder: AssertionBuilder,
message: str,
):
"""Add an assertion bound to the given expression
Args:
bound_expr (p.Expr): the expression before which the assertion
must be generated
inputs (list[p.Expr]): the list of input expressions (see :class:`Assertion`)
builder (AssertionBuilder): the builder callback (see :class:`Assertion`)
message (str): the assertion message
"""
self.assertions.append(
Assertion(
bound_expr=bound_expr,
@@ -40,20 +74,51 @@ class AssertionCollector:
)
def remove(self, assertion: Assertion):
"""Remove the given assertion from the collection
Args:
assertion (Assertion): the assertion to remove
"""
try:
self.assertions.remove(assertion)
except ValueError:
pass
def define(self, name: str, stmt: ast.stmt):
"""Register a statement definition
This method will only register the first definition of any given name
Args:
name (str): the name of the definition
stmt (ast.stmt): the definition statement, like a function def
"""
if name not in self.definitions:
self.definitions[name] = stmt
def get_definitions(self) -> list[ast.stmt]:
"""Get the list of definitions
Returns:
list[ast.stmt]: the list of definitions
"""
return list(self.definitions.values())
def get_assertions(self) -> list[Assertion]:
"""Get the list of assertions
Returns:
list[Assertion]: the list of assertions
"""
return self.assertions
def get_assertions_for(self, expr: p.Expr) -> list[Assertion]:
"""Get the list of assertions bound to a given expression
Args:
expr (p.Expr): the expression
Returns:
list[Assertion]: the list of assertions bound to `expr`
"""
return list(filter(lambda a: a.is_bound_to(expr), self.assertions))

79
midas/generator/constraints.py
View File

@@ -40,6 +40,8 @@ COMPARISON_OPERATORS: dict[TokenType, type[ast.cmpop]] = {
class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
"""Class to generate Python code for constraint expressions"""
def __init__(self, types: TypesRegistry):
self.types: TypesRegistry = types
self._id: int = 0
@@ -47,9 +49,22 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
self._aliases: dict[str, str] = {}
def get_definitions(self) -> list[ast.stmt]:
"""Get the list of definitions
Returns:
list[ast.stmt]: the list of definitions
"""
return self._definitions
def generate(self, expr: m.Expr) -> ast.expr:
"""Translate the given Midas expression to a Python expression
Args:
expr (m.Expr): the expression to translate
Returns:
ast.expr: the equivalent Python expression
"""
match expr:
case m.VariableExpr():
return expr.accept(self)
@@ -75,6 +90,14 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
return ast.Name(id=alias)
def make_alias(self, name: Optional[str]) -> str:
"""Get a unique alias for a predicate
Args:
name (Optional[str]): the name of the predicate as defined by the user
Returns:
str: a unique name
"""
suffix: str
if name is None:
suffix = f"p{self._id}"
@@ -85,6 +108,15 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
return alias
def make_definition(self, name: str, predicate: Predicate) -> ast.stmt:
"""Translate the given predicate to a Python definition (or assignment)
Args:
name (str): the name of the predicate
predicate (Predicate): the predicate
Returns:
ast.stmt: the equivalent Python statement
"""
body: ast.expr = predicate.body.accept(self)
if predicate.alias:
return ast.Assign(
@@ -96,6 +128,14 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
return self.make_func(name, [ast.Return(value=body)], predicate.type)
def make_args(self, params: ParamSpec) -> ast.arguments:
"""Translate the given parameter spec into an `ast.arguments` node
Args:
params (ParamSpec): the parameter spec to translate
Returns:
ast.arguments: the equivalent `ast.arguments`
"""
return ast.arguments(
posonlyargs=[
ast.arg(
@@ -125,6 +165,33 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
def make_func(
self, name: str, inner_body: list[ast.stmt], type: Type, level: int = 0
) -> ast.stmt:
"""Generate a Python function def with the given name, body and signature
If `type` returns a function, the curried arguments are separated into
inner methods.
For example, if `type` is `(a: int) -> (b: int) -> (c: int) -> int`, the
following function would be generated:
```python
def predicate(a: int):
def inner0(b: int):
def inner1(c: int):
return ...
return inner1
return inner0
```
Args:
name (str): the name of the outer function
inner_body (list[ast.stmt]): the body of the innermost function
type (Type): the function type / signature
level (int, optional): the current nesting level. Defaults to 0.
Raises:
ValueError: if `type` is not a function
Returns:
ast.stmt: the equivalent Python function definition
"""
match type:
case Function(params=params, returns=Function()):
inner_name: str = f"inner{level}"
@@ -152,6 +219,18 @@ class ConstraintGenerator(m.Expr.Visitor[ast.expr]):
raise ValueError(f"Expected function, got {type!r}")
def get_predicate(self, name: str) -> Optional[ast.expr]:
"""Get a predicate's alias, and generate its definition if first reference
When calling this function for the first time for a given predicate,
a Python definition and an alias are generated. Subsequent calls only
return the alias, without re-generating the predicate's definition
Args:
name (str): the predicate's name
Returns:
Optional[ast.expr]: the predicate's alias, or `None` if it is not defined
"""
if name not in self._aliases:
predicate: Optional[Predicate] = self.types.lookup_predicate(name)
if predicate is None:

247
midas/generator/generator.py
View File

@@ -40,11 +40,23 @@ from midas.utils import TypedAST
@dataclass
class Scope:
"""A simple structure to store assertions an aliases defined in a scope"""
pre_assertions: list[ast.stmt] = field(default_factory=list[ast.stmt])
"""A list of assertions that must be generated before the scope"""
aliases: list[str] = field(default_factory=list[str])
"""A list of aliases defined in the scope, that can be discard afterwards"""
class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
"""
A class to translate the custom Python AST back into raw `ast` nodes
This class is also responsible for generating assertions, functions for
predicates and other code necessary to ensure runtime safety.
"""
IS_DATAFRAME_FUNC = "__midas_is_dataframe__"
IS_COLUMN_FUNC = "__midas_is_column__"
@@ -72,9 +84,22 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
self.define_is_column: bool = False
def set_src_path(self, path: Path):
"""Set the current source file path
Args:
path (Path): the new source file path
"""
self.rel_src_path = path.resolve().relative_to(self.workdir)
def generate_ast(self, typed_ast: TypedAST) -> ast.AST:
"""Translate the given type checked AST into a Python `ast.AST`
Args:
typed_ast (TypedAST): the type checked Python AST
Returns:
ast.AST: the generated raw AST
"""
self._typed_ast = typed_ast
body: list[ast.stmt] = self._visit_body(typed_ast.stmts, can_be_empty=True)
predicates: list[ast.stmt] = self._constraint_generator.get_definitions()
@@ -103,6 +128,29 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
out_path: Optional[Path] = None,
type_files: Optional[list[tuple[Path, Optional[str]]]] = None,
) -> Path:
"""Generate all project files for the given source file and AST
This function calls :func:`generate_ast` to generate the output AST,
unparses it to runnable Python code, and also generates stubs for
user-defined Midas types in the same output directory
Args:
typed_ast (TypedAST): the type-checked AST
src_path (Path): the source file path
out_path (Optional[Path], optional): the output file path. If `None`,
the relative path of the source file to the working directory is
used to compute an equivalent path in the build directory.
Defaults to None.
type_files (Optional[list[tuple[Path, Optional[str]]]], optional):
the list of Midas files used to type check the AST. Defaults to None.
Raises:
ValueError: if `out_path` is `None` and the computed path is outside
the build directory
Returns:
Path: the actual `out_path` used
"""
self.set_src_path(src_path)
if out_path is None:
if self.build_dir.exists():
@@ -131,6 +179,12 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
return out_path
def generate_stubs(self, in_path: Path, out_path: Path):
"""Generate stubs from the given Midas file
Args:
in_path (Path): the Midas file path
out_path (Path): the stubs output file path
"""
checker = TypeChecker()
checker.import_midas(in_path)
generator = StubsGenerator(checker.types)
@@ -140,6 +194,18 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
out_path.write_text(output)
def convert(self, expr: p.Expr) -> ast.expr:
"""Translate an expression
If the expression already has an alias, it is returned.
If assertions are defined for the given expression (in :attr:`TypedAST.assertions`),
they are materialized and added to the current scope.
Args:
expr (p.Expr): the expression to translate
Returns:
ast.expr: the translated expression
"""
for expr2, alias in self._aliases:
if expr2 == expr:
return alias
@@ -256,6 +322,14 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
)
def make_args(self, params: p.ParamSpec) -> ast.arguments:
"""Translate a parameter spec into an `ast.arguments` node
Args:
params (p.ParamSpec): the parameter spec
Returns:
ast.arguments: the equivalent `ast.arguments`
"""
return ast.arguments(
posonlyargs=[ast.arg(arg=param.name) for param in params.pos],
args=[ast.arg(arg=param.name) for param in params.mixed],
@@ -325,6 +399,14 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
)
def _convert_imports(self, imports: list[p.ImportAlias]) -> list[ast.alias]:
"""Translate a list of import aliases
Args:
imports (list[p.ImportAlias]): the import aliases to translate
Returns:
list[ast.alias]: the translated aliases
"""
return [
ast.alias(
name=import_.name,
@@ -339,6 +421,21 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def _visit_body(
self, stmts: list[p.Stmt], can_be_empty: bool = False
) -> list[ast.stmt]:
"""Translate a list of statements
Assertions generated while translating a statement are inserted before it,
and aliases are deleted after the statement they're used in.
Extraneous `pass` statements are automatically removed
Args:
stmts (list[p.Stmt]): the statements to translate
can_be_empty (bool, optional): if `False` and no statement is
generated, an `ast.Pass` statement is returned. Defaults to False.
Returns:
list[ast.stmt]: the generated statements
"""
generated: list[ast.stmt] = []
for stmt in stmts:
scope = Scope()
@@ -361,6 +458,20 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
return generated
def _make_alias(self, node: p.Expr, expr: ast.expr) -> ast.expr:
"""Generate a unique alias for the given expression
This function creates a unique name, generates an assignment statement
to define the alias before the current statement, adds the alias to the
list of aliases defined in the current statement, and returns an
expression that can be used in place of `expr`
Args:
node (p.Expr): the AST node that generated `expr`
expr (ast.expr): the expression to alias
Returns:
ast.expr: the generated alias reference
"""
name: str = f"__midas_a{self._alias_count}__"
alias = ast.Name(id=name)
self._alias_count += 1
@@ -375,6 +486,15 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
return alias
def _build_assert(self, expr: ast.expr, message: str | ast.expr) -> ast.stmt:
"""Build an assert statement from the given test expression and message
Args:
expr (ast.expr): the test expression
message (str | ast.expr): the assert message
Returns:
ast.stmt: the assert statement
"""
if isinstance(message, str):
message = ast.Constant(value=message)
return ast.Assert(
@@ -383,9 +503,25 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
)
def _add_assert(self, assertion: ast.stmt):
"""Append the given assertion to the current scope
Args:
assertion (ast.stmt): the assertion to add
"""
self._scopes[-1].pre_assertions.append(assertion)
def _get_expr_type(self, query: p.Expr) -> Type:
"""Get the type of the given expression as computed by the type checker
Args:
query (p.Expr): the expression
Raises:
RuntimeError: if no type judgment can be found for `query`
Returns:
Type: the type of `expr`
"""
for expr, type in self._typed_ast.judgements:
if expr == query:
return type
@@ -394,6 +530,17 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def _make_cast_asserts(
self, src_location: Location, expr: ast.expr, type: Type
) -> list[ast.stmt]:
"""Generate assertions for the given cast expression
Args:
src_location (Location): the location of the cast expression in
the source file
expr (ast.expr): the expression being cast
type (Type): the target type
Returns:
list[ast.stmt]: the generated assertion statements
"""
match type:
case UnknownType() | TopType():
return []
@@ -548,6 +695,24 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
type: Type,
extra: Optional[str] = None,
) -> ast.expr:
"""Build an AST node for a cast assertion message
The generated Python code looks like:
```python
f"file.py:L1:1: CastError: Cannot cast {type(expr).__name__} to Type"
```
Args:
location (Location): the location of the cast expression in the
source file
expr (ast.expr): the expression being cast
type (Type): the target type
extra (Optional[str], optional): extra text to append at the end of
the message. Defaults to None.
Returns:
ast.expr: the generated message (as an f-string)
"""
loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
# f"file.py:L1:1: CastError: Cannot cast {type(expr).__name__} to Type"
return ast.JoinedStr(
@@ -571,6 +736,17 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def _make_constraint_assert(
self, src_location: Location, expr: ast.expr, constraint: m.Expr
) -> ast.stmt:
"""Build an assertion for the given constraint on the given expression
Args:
src_location (Location): the location of the cast expression in the
source file
expr (ast.expr): the expression subject to `constraint`
constraint (m.Expr): the constraint applied on `expr`
Returns:
ast.stmt: the assert statement checking the constraint
"""
test_func: ast.expr = self._get_constraint(constraint)
return self._build_assert(
ast.Call(
@@ -578,12 +754,22 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
args=[expr],
keywords=[],
),
self._make_constraint_assert_message(src_location, expr, constraint),
self._make_constraint_assert_message(src_location, constraint),
)
def _make_constraint_assert_message(
self, location: Location, expr: ast.expr, constraint: m.Expr
self, location: Location, constraint: m.Expr
) -> ast.expr:
"""Build an assert message for the given constraint
Args:
location (Location): the location of the cast expression in the
source file
constraint (m.Expr): the constraint
Returns:
ast.expr: the assert message
"""
printer = MidasPrinter()
constraint_str: str = printer.print(constraint)
loc_str: str = f"{self.rel_src_path}:L{location.lineno}:{location.col_offset+1}"
@@ -593,6 +779,14 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
)
def _get_constraint(self, expr: m.Expr) -> ast.expr:
"""Get or generate a Python expression for the given constraint
Args:
expr (m.Expr): the constraint
Returns:
ast.expr: an equivalent Python expression
"""
for expr2, constraint in self._constraints:
if expr2 == expr:
return constraint
@@ -602,10 +796,18 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
return constraint
def _is_dataframe_definition(self) -> ast.stmt:
"""
"""Build a function def to check if a value is a dataframe
The function is defined as:
```python
def IS_DATAFRAME_FUNC(obj) -> bool:
import pandas as pd
return isinstance(obj, pd.DataFrame)
```
where `IS_DATAFRAME_FUNC` is replaced by :attr:`IS_DATAFRAME_FUNC`
Returns:
ast.stmt: the function def
"""
return ast.FunctionDef(
@@ -638,10 +840,18 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
)
def _is_column_definition(self) -> ast.stmt:
"""
"""Build a function def to check if a value is a column
The function is defined as:
```python
def IS_COLUMN_FUNC(obj) -> bool:
import pandas as pd
return isinstance(obj, pd.Series)
```
where `IS_COLUMN_FUNC` is replaced by :attr:`IS_COLUMN_FUNC`
Returns:
ast.stmt: the function def
"""
return ast.FunctionDef(
@@ -676,6 +886,18 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
def _make_column_inner_assert(
self, src_location: Location, column: ast.expr, type: ColumnType
) -> Optional[ast.stmt]:
"""Build a for-loop checking the type of values inside a column
Args:
src_location (Location): the location of the cast expression in the
source file
column (ast.expr): the column being cast
type (ColumnType): the type of the column
Returns:
Optional[ast.stmt]: a for-loop checking the values, or `None` if no
assertions are necessary
"""
# TODO: improve message, maybe chain contexts
col: ast.expr = ast.Name(id="col")
body: list[ast.stmt] = self._make_cast_asserts(src_location, col, type.type)
@@ -689,6 +911,14 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
)
def _convert_assertion(self, assertion: Assertion) -> ast.stmt:
"""Generate a Python assert statement for the given assertion
Args:
assertion (Assertion): the assertion to translate
Returns:
ast.stmt: the generated assert statement
"""
inputs: list[ast.expr] = []
for input in assertion.inputs:
@@ -704,6 +934,15 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
)
def _apply_assertions(self, expr: p.Expr, assertions: list[Assertion]) -> ast.expr:
"""Translate the given expression, adding linked assertions to the scope
Args:
expr (p.Expr): the expression to translate
assertions (list[Assertion]): the list of assertions linked to `expr`
Returns:
ast.expr: the translated expression
"""
for assertion in assertions:
assert_stmt: ast.stmt
assert_stmt = self._convert_assertion(assertion)

112
midas/generator/stubs.py
View File

@@ -32,6 +32,8 @@ Empty = ast.Constant(value=...)
class StubsGenerator:
"""A class to generate Python stubs for user-defined Midas types"""
def __init__(self, types: TypesRegistry) -> None:
self.types: TypesRegistry = types
self.stubs: list[ast.stmt] = []
@@ -43,6 +45,11 @@ class StubsGenerator:
self.substitutions: dict[str, dict[str, Type]] = {}
def generate_stubs(self) -> ast.Module:
"""Generate a Python module of stubs for all user-defined types
Returns:
ast.Module: the generated module
"""
self.stubs = []
self.typing_imports = set()
self.import_pandas = False
@@ -92,6 +99,12 @@ class StubsGenerator:
return ast.Module(body=imports + self.stubs, type_ignores=[])
def generate_stub(self, name: str, type: Type):
"""Generate a stub for the given type
Args:
name (str): the name of the type
type (Type): the type
"""
base_type: Type = type
# TODO: improve
@@ -129,6 +142,17 @@ class StubsGenerator:
self.add_stub(stub)
def get_bases(self, type: Type) -> tuple[list[ast.expr], dict[str, Type]]:
"""Get the list of class bases and type parameter substitutions for a type
Args:
type (Type): the type whose bases to get
Returns:
tuple[list[ast.expr], dict[str, Type]]: a tuple containing the list
of class bases (already translated to Python AST nodes), and a
mapping of type parameter substitutions (to replace them with
their generated aliases)
"""
match type:
case DerivedType(type=base):
return [self.dump_type(base)], {}
@@ -173,6 +197,16 @@ class StubsGenerator:
def generate_body(
self, members: dict[str, Member], substitutions: dict[str, Type]
) -> list[ast.stmt]:
"""Generate a class body given its members
Args:
members (dict[str, Member]): the class members
substitutions (dict[str, Type]): a mapping of type parameter
substitutions (to replace them with their generated aliases)
Returns:
list[ast.stmt]: the generated class body statements
"""
if len(members) == 0:
return [ast.Expr(value=Empty)]
@@ -194,6 +228,14 @@ class StubsGenerator:
return body
def dump_type(self, type: Type) -> ast.expr:
"""Translate a type to a Python expression
Args:
type (Type): the type to translate
Returns:
ast.expr: the generated Python expression
"""
match type:
case DerivedType(name=name) | GenericType(name=name) if (
name in self.substitutions
@@ -319,6 +361,17 @@ class StubsGenerator:
def dump_method(
self, name: str, method: Type, overloaded: bool = False
) -> list[ast.stmt]:
"""Generate definitions for a method
Args:
name (str): the method's name
method (Type): the method's type
overloaded (bool, optional): whether this method is part of an
overloaded method (used when called recursively). Defaults to False.
Returns:
list[ast.stmt]: the generated function definitions
"""
match method:
case Function():
if overloaded:
@@ -347,6 +400,16 @@ class StubsGenerator:
]
def dump_params(self, params: ParamSpec, with_self: bool = False) -> ast.arguments:
"""Generate an `ast.arguments` node for the given parameter spec
Args:
params (ParamSpec): the parameter spec to translate
with_self (bool, optional): whether to include a `self` parameter.
Defaults to False.
Returns:
ast.arguments: the generate Python AST node
"""
pos: list[ast.arg] = [
ast.arg(
arg=f"_{param.pos}",
@@ -389,6 +452,14 @@ class StubsGenerator:
)
def define_protocol(self, func: Function) -> str:
"""Generate a :class:`Protocol` to use in a function stub
Args:
func (Function): the function signature to define
Returns:
str: the name of the generated protocol
"""
self.add_typing_import("Protocol")
name: str = self.new_protocol_name()
protocol = ast.ClassDef(
@@ -410,33 +481,74 @@ class StubsGenerator:
return name
def new_protocol_name(self) -> str:
"""Get a unique protocol name
Returns:
str: the unique protocol name
"""
name: str = f"_Protocol{self.protocol_idx}"
self.protocol_idx += 1
return name
def new_stub_name(self) -> str:
"""Get a unique stub name
Returns:
str: the unique stub name
"""
name: str = f"_Stub_{self.stub_idx}"
self.stub_idx += 1
return name
def new_type_var_name(self) -> str:
"""Get a unique type variable name
Returns:
str: the unique type variable name
"""
name: str = f"_T{self.type_var_idx}"
self.type_var_idx += 1
return name
def add_stub(self, stub: ast.stmt):
"""Append the given statement to the output
Args:
stub (ast.stmt): the statement to append
"""
self.stubs.append(stub)
def add_typing_import(self, name: str):
"""Add the given name to the list of names to import from `typing`
Args:
name (str): the name to import
"""
self.typing_imports.add(name)
def define_type_vars(self, vars: list[TypeVar]) -> list[TypeVar]:
"""Define aliases for the given type variables
Args:
vars (list[TypeVar]): the variables to define
Returns:
list[TypeVar]: new type variables named with the generated aliases
"""
vars2: list[TypeVar] = []
for var in vars:
vars2.append(self.define_type_var(var))
return vars2
def define_type_var(self, var: TypeVar) -> TypeVar:
"""Define a type variable alias
Args:
var (TypeVar): the type variable to define
Returns:
TypeVar: a new type variable named with a uniquely generated alias
"""
name: str = self.new_type_var_name()
self.add_typing_import("TypeVar")