from __future__ import annotations import ast import builtins import typing from ast import ( AST, Add, AnnAssign, Assign, AsyncFunctionDef, Attribute, AugAssign, BinOp, BitAnd, BitOr, BitXor, Call, ClassDef, Constant, Dict, Div, Expr, Expression, FloorDiv, FunctionDef, If, Import, ImportFrom, List, Load, LShift, MatMult, Mod, Module, Mult, Name, NamedExpr, NodeTransformer, NodeVisitor, Pass, Pow, Return, RShift, Starred, Store, Sub, Subscript, Tuple, Yield, YieldFrom, alias, copy_location, expr, fix_missing_locations, keyword, walk, ) from collections import defaultdict from collections.abc import Generator, Sequence from contextlib import contextmanager from copy import deepcopy from dataclasses import dataclass, field from typing import Any, ClassVar, cast, overload from typing import __all__ as typing_all from typing_extensions import __all__ as typing_extensions_all PreliminaryNameTypePair: typing.TypeAlias = tuple[Constant, "expr | None"] NameTypePair: typing.TypeAlias = tuple[Constant, expr] generator_names = ( "typing.Generator", "collections.abc.Generator", "typing.Iterator", "collections.abc.Iterator", "typing.Iterable", "collections.abc.Iterable", "typing.AsyncIterator", "collections.abc.AsyncIterator", "typing.AsyncIterable", "collections.abc.AsyncIterable", "typing.AsyncGenerator", "collections.abc.AsyncGenerator", ) anytype_names = ( "typing.Any", "typing_extensions.Any", ) literal_names = ( "typing.Literal", "typing_extensions.Literal", ) annotated_names = ( "typing.Annotated", "typing_extensions.Annotated", ) ignore_decorators = ( "typing.no_type_check", "typeguard.typeguard_ignore", ) aug_assign_functions = { Add: "iadd", Sub: "isub", Mult: "imul", MatMult: "imatmul", Div: "itruediv", FloorDiv: "ifloordiv", Mod: "imod", Pow: "ipow", LShift: "ilshift", RShift: "irshift", BitAnd: "iand", BitXor: "ixor", BitOr: "ior", } @dataclass class TransformMemo: node: Module | ClassDef | FunctionDef | AsyncFunctionDef | None parent: TransformMemo | None path: tuple[str, ...] joined_path: Constant = field(init=False) return_annotation: expr | None = None yield_annotation: expr | None = None send_annotation: expr | None = None is_async: bool = False local_names: set[str] = field(init=False, default_factory=set) imported_names: dict[str, str] = field(init=False, default_factory=dict) ignored_names: set[str] = field(init=False, default_factory=set) load_names: defaultdict[str, dict[str, Name]] = field( init=False, default_factory=lambda: defaultdict(dict) ) has_yield_expressions: bool = field(init=False, default=False) has_return_expressions: bool = field(init=False, default=False) memo_var_name: Name | None = field(init=False, default=None) should_instrument: bool = field(init=False, default=True) variable_annotations: dict[str, expr] = field(init=False, default_factory=dict) configuration_overrides: dict[str, Any] = field(init=False, default_factory=dict) code_inject_index: int = field(init=False, default=0) def __post_init__(self) -> None: elements: list[str] = [] memo = self while isinstance(memo.node, (ClassDef, FunctionDef, AsyncFunctionDef)): elements.insert(0, memo.node.name) if not memo.parent: break memo = memo.parent if isinstance(memo.node, (FunctionDef, AsyncFunctionDef)): elements.insert(0, "") self.joined_path = Constant(".".join(elements)) # Figure out where to insert instrumentation code if self.node: for index, child in enumerate(self.node.body): if isinstance(child, ImportFrom) and child.module == "__future__": # (module only) __future__ imports must come first continue elif ( isinstance(child, Expr) and isinstance(child.value, Constant) and isinstance(child.value.value, str) ): continue # docstring self.code_inject_index = index break def get_unused_name(self, name: str) -> str: memo: TransformMemo | None = self while memo is not None: if name in memo.local_names: memo = self name += "_" else: memo = memo.parent self.local_names.add(name) return name def is_ignored_name(self, expression: expr | Expr | None) -> bool: top_expression = ( expression.value if isinstance(expression, Expr) else expression ) if isinstance(top_expression, Attribute) and isinstance( top_expression.value, Name ): name = top_expression.value.id elif isinstance(top_expression, Name): name = top_expression.id else: return False memo: TransformMemo | None = self while memo is not None: if name in memo.ignored_names: return True memo = memo.parent return False def get_memo_name(self) -> Name: if not self.memo_var_name: self.memo_var_name = Name(id="memo", ctx=Load()) return self.memo_var_name def get_import(self, module: str, name: str) -> Name: if module in self.load_names and name in self.load_names[module]: return self.load_names[module][name] qualified_name = f"{module}.{name}" if name in self.imported_names and self.imported_names[name] == qualified_name: return Name(id=name, ctx=Load()) alias = self.get_unused_name(name) node = self.load_names[module][name] = Name(id=alias, ctx=Load()) self.imported_names[name] = qualified_name return node def insert_imports(self, node: Module | FunctionDef | AsyncFunctionDef) -> None: """Insert imports needed by injected code.""" if not self.load_names: return # Insert imports after any "from __future__ ..." imports and any docstring for modulename, names in self.load_names.items(): aliases = [ alias(orig_name, new_name.id if orig_name != new_name.id else None) for orig_name, new_name in sorted(names.items()) ] node.body.insert(self.code_inject_index, ImportFrom(modulename, aliases, 0)) def name_matches(self, expression: expr | Expr | None, *names: str) -> bool: if expression is None: return False path: list[str] = [] top_expression = ( expression.value if isinstance(expression, Expr) else expression ) if isinstance(top_expression, Subscript): top_expression = top_expression.value elif isinstance(top_expression, Call): top_expression = top_expression.func while isinstance(top_expression, Attribute): path.insert(0, top_expression.attr) top_expression = top_expression.value if not isinstance(top_expression, Name): return False if top_expression.id in self.imported_names: translated = self.imported_names[top_expression.id] elif hasattr(builtins, top_expression.id): translated = "builtins." + top_expression.id else: translated = top_expression.id path.insert(0, translated) joined_path = ".".join(path) if joined_path in names: return True elif self.parent: return self.parent.name_matches(expression, *names) else: return False def get_config_keywords(self) -> list[keyword]: if self.parent and isinstance(self.parent.node, ClassDef): overrides = self.parent.configuration_overrides.copy() else: overrides = {} overrides.update(self.configuration_overrides) return [keyword(key, value) for key, value in overrides.items()] class NameCollector(NodeVisitor): def __init__(self) -> None: self.names: set[str] = set() def visit_Import(self, node: Import) -> None: for name in node.names: self.names.add(name.asname or name.name) def visit_ImportFrom(self, node: ImportFrom) -> None: for name in node.names: self.names.add(name.asname or name.name) def visit_Assign(self, node: Assign) -> None: for target in node.targets: if isinstance(target, Name): self.names.add(target.id) def visit_NamedExpr(self, node: NamedExpr) -> Any: if isinstance(node.target, Name): self.names.add(node.target.id) def visit_FunctionDef(self, node: FunctionDef) -> None: pass def visit_ClassDef(self, node: ClassDef) -> None: pass class GeneratorDetector(NodeVisitor): """Detects if a function node is a generator function.""" contains_yields: bool = False in_root_function: bool = False def visit_Yield(self, node: Yield) -> Any: self.contains_yields = True def visit_YieldFrom(self, node: YieldFrom) -> Any: self.contains_yields = True def visit_ClassDef(self, node: ClassDef) -> Any: pass def visit_FunctionDef(self, node: FunctionDef | AsyncFunctionDef) -> Any: if not self.in_root_function: self.in_root_function = True self.generic_visit(node) self.in_root_function = False def visit_AsyncFunctionDef(self, node: AsyncFunctionDef) -> Any: self.visit_FunctionDef(node) class AnnotationTransformer(NodeTransformer): type_substitutions: ClassVar[dict[str, tuple[str, str]]] = { "builtins.dict": ("typing", "Dict"), "builtins.list": ("typing", "List"), "builtins.tuple": ("typing", "Tuple"), "builtins.set": ("typing", "Set"), "builtins.frozenset": ("typing", "FrozenSet"), } def __init__(self, transformer: TypeguardTransformer): self.transformer = transformer self._memo = transformer._memo self._level = 0 def visit(self, node: AST) -> Any: # Don't process Literals if isinstance(node, expr) and self._memo.name_matches(node, *literal_names): return node self._level += 1 new_node = super().visit(node) self._level -= 1 if isinstance(new_node, Expression) and not hasattr(new_node, "body"): return None # Return None if this new node matches a variation of typing.Any if ( self._level == 0 and isinstance(new_node, expr) and self._memo.name_matches(new_node, *anytype_names) ): return None return new_node def visit_BinOp(self, node: BinOp) -> Any: self.generic_visit(node) if isinstance(node.op, BitOr): # If either branch of the BinOp has been transformed to `None`, it means # that a type in the union was ignored, so the entire annotation should be # ignored if not hasattr(node, "left") or not hasattr(node, "right"): return None # Return Any if either side is Any if self._memo.name_matches(node.left, *anytype_names): return node.left elif self._memo.name_matches(node.right, *anytype_names): return node.right return node def visit_Attribute(self, node: Attribute) -> Any: if self._memo.is_ignored_name(node): return None return node def visit_Subscript(self, node: Subscript) -> Any: if self._memo.is_ignored_name(node.value): return None # The subscript of typing(_extensions).Literal can be any arbitrary string, so # don't try to evaluate it as code if node.slice: if isinstance(node.slice, Tuple): if self._memo.name_matches(node.value, *annotated_names): # Only treat the first argument to typing.Annotated as a potential # forward reference items = cast( typing.List[expr], [self.visit(node.slice.elts[0])] + node.slice.elts[1:], ) else: items = cast( typing.List[expr], [self.visit(item) for item in node.slice.elts], ) # If this is a Union and any of the items is Any, erase the entire # annotation if self._memo.name_matches(node.value, "typing.Union") and any( item is None or ( isinstance(item, expr) and self._memo.name_matches(item, *anytype_names) ) for item in items ): return None # If all items in the subscript were Any, erase the subscript entirely # (but not for an empty subscript like tuple[()], which is meaningful) if items and all(item is None for item in items): return node.value for index, item in enumerate(items): if item is None: items[index] = self.transformer._get_import("typing", "Any") node.slice.elts = items else: self.generic_visit(node) # If the transformer erased the slice entirely, just return the node # value without the subscript (unless it's Optional, in which case erase # the node entirely if self._memo.name_matches( node.value, "typing.Optional" ) and not hasattr(node, "slice"): return None if not hasattr(node, "slice"): return node.value return node def visit_Name(self, node: Name) -> Any: if self._memo.is_ignored_name(node): return None return node def visit_Call(self, node: Call) -> Any: # Don't recurse into calls return node def visit_Constant(self, node: Constant) -> Any: if isinstance(node.value, str): expression = ast.parse(node.value, mode="eval") new_node = self.visit(expression) if new_node: return copy_location(new_node.body, node) else: return None return node class TypeguardTransformer(NodeTransformer): def __init__( self, target_path: Sequence[str] | None = None, target_lineno: int | None = None ) -> None: self._target_path = tuple(target_path) if target_path else None self._memo = self._module_memo = TransformMemo(None, None, ()) self.names_used_in_annotations: set[str] = set() self.target_node: FunctionDef | AsyncFunctionDef | None = None self.target_lineno = target_lineno def generic_visit(self, node: AST) -> AST: has_non_empty_body_initially = bool(getattr(node, "body", None)) initial_type = type(node) node = super().generic_visit(node) if ( type(node) is initial_type and has_non_empty_body_initially and hasattr(node, "body") and not node.body ): # If we have still the same node type after transformation # but we've optimised it's body away, we add a `pass` statement. node.body = [Pass()] return node @contextmanager def _use_memo( self, node: ClassDef | FunctionDef | AsyncFunctionDef ) -> Generator[None, Any, None]: new_memo = TransformMemo(node, self._memo, self._memo.path + (node.name,)) old_memo = self._memo self._memo = new_memo if isinstance(node, (FunctionDef, AsyncFunctionDef)): new_memo.should_instrument = ( self._target_path is None or new_memo.path == self._target_path ) if new_memo.should_instrument: # Check if the function is a generator function detector = GeneratorDetector() detector.visit(node) # Extract yield, send and return types where possible from a subscripted # annotation like Generator[int, str, bool] return_annotation = deepcopy(node.returns) if detector.contains_yields and new_memo.name_matches( return_annotation, *generator_names ): if isinstance(return_annotation, Subscript): if isinstance(return_annotation.slice, Tuple): items = return_annotation.slice.elts else: items = [return_annotation.slice] if len(items) > 0: new_memo.yield_annotation = self._convert_annotation( items[0] ) if len(items) > 1: new_memo.send_annotation = self._convert_annotation( items[1] ) if len(items) > 2: new_memo.return_annotation = self._convert_annotation( items[2] ) else: new_memo.return_annotation = self._convert_annotation( return_annotation ) if isinstance(node, AsyncFunctionDef): new_memo.is_async = True yield self._memo = old_memo def _get_import(self, module: str, name: str) -> Name: memo = self._memo if self._target_path else self._module_memo return memo.get_import(module, name) @overload def _convert_annotation(self, annotation: None) -> None: ... @overload def _convert_annotation(self, annotation: expr) -> expr: ... def _convert_annotation(self, annotation: expr | None) -> expr | None: if annotation is None: return None # Convert PEP 604 unions (x | y) and generic built-in collections where # necessary, and undo forward references new_annotation = cast(expr, AnnotationTransformer(self).visit(annotation)) if isinstance(new_annotation, expr): new_annotation = ast.copy_location(new_annotation, annotation) # Store names used in the annotation names = {node.id for node in walk(new_annotation) if isinstance(node, Name)} self.names_used_in_annotations.update(names) return new_annotation def visit_Name(self, node: Name) -> Name: self._memo.local_names.add(node.id) return node def visit_Module(self, node: Module) -> Module: self._module_memo = self._memo = TransformMemo(node, None, ()) self.generic_visit(node) self._module_memo.insert_imports(node) fix_missing_locations(node) return node def visit_Import(self, node: Import) -> Import: for name in node.names: self._memo.local_names.add(name.asname or name.name) self._memo.imported_names[name.asname or name.name] = name.name return node def visit_ImportFrom(self, node: ImportFrom) -> ImportFrom: for name in node.names: if name.name == "*": match node.module: case "typing": imported_names = typing_all case "typing_extensions": imported_names = typing_extensions_all case _: continue for item in imported_names: self._memo.local_names.add(item) self._memo.imported_names.setdefault(item, f"{node.module}.{item}") else: alias = name.asname or name.name self._memo.local_names.add(alias) self._memo.imported_names[alias] = f"{node.module}.{name.name}" return node def visit_ClassDef(self, node: ClassDef) -> ClassDef | None: self._memo.local_names.add(node.name) # Eliminate top level classes not belonging to the target path if ( self._target_path is not None and not self._memo.path and node.name != self._target_path[0] ): return None with self._use_memo(node): for decorator in node.decorator_list.copy(): if self._memo.name_matches(decorator, "typeguard.typechecked"): # Remove the decorator to prevent duplicate instrumentation node.decorator_list.remove(decorator) # Store any configuration overrides if isinstance(decorator, Call) and decorator.keywords: self._memo.configuration_overrides.update( {kw.arg: kw.value for kw in decorator.keywords if kw.arg} ) self.generic_visit(node) return node def visit_FunctionDef( self, node: FunctionDef | AsyncFunctionDef ) -> FunctionDef | AsyncFunctionDef | None: """ Injects type checks for function arguments, and for a return of None if the function is annotated to return something else than Any or None, and the body ends without an explicit "return". """ self._memo.local_names.add(node.name) # Eliminate top level functions not belonging to the target path if ( self._target_path is not None and not self._memo.path and node.name != self._target_path[0] ): return None # Skip instrumentation if we're instrumenting the whole module and the function # contains either @no_type_check or @typeguard_ignore if self._target_path is None: for decorator in node.decorator_list: if self._memo.name_matches(decorator, *ignore_decorators): return node with self._use_memo(node): arg_annotations: dict[str, Any] = {} if self._target_path is None or self._memo.path == self._target_path: # Find line number we're supposed to match against if node.decorator_list: first_lineno = node.decorator_list[0].lineno else: first_lineno = node.lineno for decorator in node.decorator_list.copy(): if self._memo.name_matches(decorator, "typing.overload"): # Remove overloads entirely return None elif self._memo.name_matches(decorator, "typeguard.typechecked"): # Remove the decorator to prevent duplicate instrumentation node.decorator_list.remove(decorator) # Store any configuration overrides if isinstance(decorator, Call) and decorator.keywords: self._memo.configuration_overrides = { kw.arg: kw.value for kw in decorator.keywords if kw.arg } if self.target_lineno == first_lineno: assert self.target_node is None self.target_node = node if node.decorator_list: self.target_lineno = node.decorator_list[0].lineno else: self.target_lineno = node.lineno all_args = node.args.posonlyargs + node.args.args + node.args.kwonlyargs # Ensure that any type shadowed by the positional or keyword-only # argument names are ignored in this function for arg in all_args: self._memo.ignored_names.add(arg.arg) # Ensure that any type shadowed by the variable positional argument name # (e.g. "args" in *args) is ignored this function if node.args.vararg: self._memo.ignored_names.add(node.args.vararg.arg) # Ensure that any type shadowed by the variable keywrod argument name # (e.g. "kwargs" in *kwargs) is ignored this function if node.args.kwarg: self._memo.ignored_names.add(node.args.kwarg.arg) for arg in all_args: annotation = self._convert_annotation(deepcopy(arg.annotation)) if annotation: arg_annotations[arg.arg] = annotation if node.args.vararg: annotation_ = self._convert_annotation(node.args.vararg.annotation) if annotation_: container = Name("tuple", ctx=Load()) subscript_slice = Tuple( [ annotation_, Constant(Ellipsis), ], ctx=Load(), ) arg_annotations[node.args.vararg.arg] = Subscript( container, subscript_slice, ctx=Load() ) if node.args.kwarg: annotation_ = self._convert_annotation(node.args.kwarg.annotation) if annotation_: container = Name("dict", ctx=Load()) subscript_slice = Tuple( [ Name("str", ctx=Load()), annotation_, ], ctx=Load(), ) arg_annotations[node.args.kwarg.arg] = Subscript( container, subscript_slice, ctx=Load() ) if arg_annotations: self._memo.variable_annotations.update(arg_annotations) self.generic_visit(node) if arg_annotations: annotations_dict = Dict( keys=[Constant(key) for key in arg_annotations.keys()], values=[ Tuple([Name(key, ctx=Load()), annotation], ctx=Load()) for key, annotation in arg_annotations.items() ], ) func_name = self._get_import( "typeguard._functions", "check_argument_types_internal" ) args = [ self._memo.joined_path, annotations_dict, self._memo.get_memo_name(), ] node.body.insert( self._memo.code_inject_index, Expr(Call(func_name, args, [])) ) # Add a checked "return None" to the end if there's no explicit return # Skip if the return annotation is None or Any if ( self._memo.return_annotation and (not self._memo.is_async or not self._memo.has_yield_expressions) and not isinstance(node.body[-1], Return) and ( not isinstance(self._memo.return_annotation, Constant) or self._memo.return_annotation.value is not None ) ): func_name = self._get_import( "typeguard._functions", "check_return_type_internal" ) return_node = Return( Call( func_name, [ self._memo.joined_path, Constant(None), self._memo.return_annotation, self._memo.get_memo_name(), ], [], ) ) # Replace a placeholder "pass" at the end if isinstance(node.body[-1], Pass): copy_location(return_node, node.body[-1]) del node.body[-1] node.body.append(return_node) # Insert code to create the call memo, if it was ever needed for this # function if self._memo.memo_var_name: memo_kwargs: dict[str, Any] = {} if self._memo.parent and isinstance(self._memo.parent.node, ClassDef): for decorator in node.decorator_list: if ( isinstance(decorator, Name) and decorator.id == "staticmethod" ): break elif ( isinstance(decorator, Name) and decorator.id == "classmethod" ): arglist = node.args.posonlyargs or node.args.args memo_kwargs["self_type"] = Name( id=arglist[0].arg, ctx=Load() ) break else: if arglist := node.args.posonlyargs or node.args.args: if node.name == "__new__": memo_kwargs["self_type"] = Name( id=arglist[0].arg, ctx=Load() ) else: memo_kwargs["self_type"] = Attribute( Name(id=arglist[0].arg, ctx=Load()), "__class__", ctx=Load(), ) # Construct the function reference # Nested functions get special treatment: the function name is added # to free variables (and the closure of the resulting function) names: list[str] = [node.name] memo = self._memo.parent while memo: if isinstance(memo.node, (FunctionDef, AsyncFunctionDef)): # This is a nested function. Use the function name as-is. del names[:-1] break elif not isinstance(memo.node, ClassDef): break names.insert(0, memo.node.name) memo = memo.parent config_keywords = self._memo.get_config_keywords() if config_keywords: memo_kwargs["config"] = Call( self._get_import("dataclasses", "replace"), [self._get_import("typeguard._config", "global_config")], config_keywords, ) self._memo.memo_var_name.id = self._memo.get_unused_name("memo") memo_store_name = Name(id=self._memo.memo_var_name.id, ctx=Store()) globals_call = Call(Name(id="globals", ctx=Load()), [], []) locals_call = Call(Name(id="locals", ctx=Load()), [], []) memo_expr = Call( self._get_import("typeguard", "TypeCheckMemo"), [globals_call, locals_call], [keyword(key, value) for key, value in memo_kwargs.items()], ) node.body.insert( self._memo.code_inject_index, Assign([memo_store_name], memo_expr), ) self._memo.insert_imports(node) # Special case the __new__() method to create a local alias from the # class name to the first argument (usually "cls") if ( isinstance(node, FunctionDef) and node.args and self._memo.parent is not None and isinstance(self._memo.parent.node, ClassDef) and node.name == "__new__" ): first_args_expr = Name(node.args.args[0].arg, ctx=Load()) cls_name = Name(self._memo.parent.node.name, ctx=Store()) node.body.insert( self._memo.code_inject_index, Assign([cls_name], first_args_expr), ) # Remove any placeholder "pass" at the end if isinstance(node.body[-1], Pass): del node.body[-1] return node def visit_AsyncFunctionDef( self, node: AsyncFunctionDef ) -> FunctionDef | AsyncFunctionDef | None: return self.visit_FunctionDef(node) def visit_Return(self, node: Return) -> Return: """This injects type checks into "return" statements.""" self.generic_visit(node) if ( self._memo.return_annotation and self._memo.should_instrument and not self._memo.is_ignored_name(self._memo.return_annotation) ): func_name = self._get_import( "typeguard._functions", "check_return_type_internal" ) old_node = node retval = old_node.value or Constant(None) node = Return( Call( func_name, [ self._memo.joined_path, retval, self._memo.return_annotation, self._memo.get_memo_name(), ], [], ) ) copy_location(node, old_node) return node def visit_Yield(self, node: Yield) -> Yield | Call: """ This injects type checks into "yield" expressions, checking both the yielded value and the value sent back to the generator, when appropriate. """ self._memo.has_yield_expressions = True self.generic_visit(node) if ( self._memo.yield_annotation and self._memo.should_instrument and not self._memo.is_ignored_name(self._memo.yield_annotation) ): func_name = self._get_import("typeguard._functions", "check_yield_type") yieldval = node.value or Constant(None) node.value = Call( func_name, [ self._memo.joined_path, yieldval, self._memo.yield_annotation, self._memo.get_memo_name(), ], [], ) if ( self._memo.send_annotation and self._memo.should_instrument and not self._memo.is_ignored_name(self._memo.send_annotation) ): func_name = self._get_import("typeguard._functions", "check_send_type") old_node = node call_node = Call( func_name, [ self._memo.joined_path, old_node, self._memo.send_annotation, self._memo.get_memo_name(), ], [], ) copy_location(call_node, old_node) return call_node return node def visit_AnnAssign(self, node: AnnAssign) -> Any: """ This injects a type check into a local variable annotation-assignment within a function body. """ self.generic_visit(node) if ( isinstance(self._memo.node, (FunctionDef, AsyncFunctionDef)) and node.annotation and isinstance(node.target, Name) ): self._memo.ignored_names.add(node.target.id) annotation = self._convert_annotation(deepcopy(node.annotation)) if annotation: self._memo.variable_annotations[node.target.id] = annotation if node.value: func_name = self._get_import( "typeguard._functions", "check_variable_assignment" ) targets_arg = List( [ Tuple( [Constant(node.target.id), annotation], ctx=Load(), ) ], ctx=Load(), ) node.value = Call( func_name, [ node.value, targets_arg, self._memo.get_memo_name(), ], [], ) return node def visit_Assign(self, node: Assign) -> Any: """ This injects a type check into a local variable assignment within a function body. The variable must have been annotated earlier in the function body. """ self.generic_visit(node) # Only instrument function-local assignments if isinstance(self._memo.node, (FunctionDef, AsyncFunctionDef)): preliminary_targets: list[ PreliminaryNameTypePair | list[PreliminaryNameTypePair] ] = [] check_required = False for node_target in node.targets: elts: Sequence[expr] if isinstance(node_target, Name): elts = [node_target] single_target = True elif isinstance(node_target, Tuple): elts = node_target.elts single_target = False else: continue annotations_: list[PreliminaryNameTypePair] = [] for exp in elts: prefix = "" if isinstance(exp, Starred): exp = exp.value prefix = "*" path: list[str] = [] while isinstance(exp, Attribute): path.insert(0, exp.attr) exp = exp.value if isinstance(exp, Name): if not path: self._memo.ignored_names.add(exp.id) path.insert(0, exp.id) name = prefix + ".".join(path) if len(path) == 1 and ( annotation := self._memo.variable_annotations.get(exp.id) ): annotations_.append((Constant(name), annotation)) check_required = True else: annotations_.append((Constant(name), None)) preliminary_targets.append( annotations_[0] if single_target else annotations_ ) if check_required: # Replace missing annotations with typing.Any targets: list[NameTypePair | list[NameTypePair]] = [] for items in preliminary_targets: if isinstance(items, list): target_list: list[NameTypePair] = [] targets.append(target_list) for key, expression_or_none in items: expression = expression_or_none or self._get_import( "typing", "Any" ) target_list.append((key, expression)) else: key, expression_or_none = items expression = expression_or_none or self._get_import( "typing", "Any" ) targets.append((key, expression)) func_name = self._get_import( "typeguard._functions", "check_variable_assignment" ) elements: list[expr] = [] for target in targets: if isinstance(target, list): elements.append( List( [ Tuple([name, ann], ctx=Load()) for name, ann in target ], ctx=Load(), ) ) else: name_constant, ann = target elements.append(Tuple([name_constant, ann], ctx=Load())) targets_arg = List(elements, ctx=Load()) node.value = Call( func_name, [node.value, targets_arg, self._memo.get_memo_name()], [], ) return node def visit_NamedExpr(self, node: NamedExpr) -> Any: """This injects a type check into an assignment expression (a := foo()).""" self.generic_visit(node) # Only instrument function-local assignments if isinstance(self._memo.node, (FunctionDef, AsyncFunctionDef)) and isinstance( node.target, Name ): self._memo.ignored_names.add(node.target.id) # Bail out if no matching annotation is found annotation = self._memo.variable_annotations.get(node.target.id) if annotation is None: return node func_name = self._get_import( "typeguard._functions", "check_variable_assignment" ) node.value = Call( func_name, [ node.value, List( [ Tuple( [Constant(node.target.id), annotation], ctx=Load(), ) ], ctx=Load(), ), self._memo.get_memo_name(), ], [], ) return node def visit_AugAssign(self, node: AugAssign) -> Any: """ This injects a type check into an augmented assignment expression (a += 1). """ self.generic_visit(node) # Only instrument function-local assignments if isinstance(self._memo.node, (FunctionDef, AsyncFunctionDef)) and isinstance( node.target, Name ): # Bail out if no matching annotation is found annotation = self._memo.variable_annotations.get(node.target.id) if annotation is None: return node # Bail out if the operator is not found (newer Python version?) try: operator_func_name = aug_assign_functions[node.op.__class__] except KeyError: return node operator_func = self._get_import("operator", operator_func_name) operator_call = Call( operator_func, [Name(node.target.id, ctx=Load()), node.value], [] ) targets_arg = List( [Tuple([Constant(node.target.id), annotation], ctx=Load())], ctx=Load(), ) check_call = Call( self._get_import("typeguard._functions", "check_variable_assignment"), [ operator_call, targets_arg, self._memo.get_memo_name(), ], [], ) return Assign(targets=[node.target], value=check_call) return node def visit_If(self, node: If) -> Any: """ This blocks names from being collected from a module-level "if typing.TYPE_CHECKING:" block, so that they won't be type checked. """ self.generic_visit(node) if ( self._memo is self._module_memo and isinstance(node.test, Name) and self._memo.name_matches(node.test, "typing.TYPE_CHECKING") ): collector = NameCollector() collector.visit(node) self._memo.ignored_names.update(collector.names) return node