From ed07b015638eba4f6f7cca29ec759ef61db7b46d Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 7 Jul 2026 12:09:44 +0200 Subject: [PATCH] docs: add docstrings to generator submodule --- midas/generator/collector.py | 65 +++++++++ midas/generator/constraints.py | 79 +++++++++++ midas/generator/generator.py | 247 ++++++++++++++++++++++++++++++++- midas/generator/stubs.py | 112 +++++++++++++++ 4 files changed, 499 insertions(+), 4 deletions(-) diff --git a/midas/generator/collector.py b/midas/generator/collector.py index fcaf545..554b721 100644 --- a/midas/generator/collector.py +++ b/midas/generator/collector.py @@ -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)) diff --git a/midas/generator/constraints.py b/midas/generator/constraints.py index 65022f2..75c9a15 100644 --- a/midas/generator/constraints.py +++ b/midas/generator/constraints.py @@ -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: diff --git a/midas/generator/generator.py b/midas/generator/generator.py index 52c2d76..35634a5 100644 --- a/midas/generator/generator.py +++ b/midas/generator/generator.py @@ -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) diff --git a/midas/generator/stubs.py b/midas/generator/stubs.py index d42f3db..f185f97 100644 --- a/midas/generator/stubs.py +++ b/midas/generator/stubs.py @@ -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")