from typing import Optional from mypy.nodes import ( AssertTypeExpr, Var, FuncItem, ClassDef, AssignmentStmt, ForStmt, WithStmt, CastExpr, TypeApplication, TypeAliasExpr, TypeVarExpr, TypedDictExpr, NamedTupleExpr, PromoteExpr, NewTypeExpr ) from mypy.types import Type from mypy.traverser import TraverserVisitor from mypy.typetraverser import TypeTraverserVisitor class MixedTraverserVisitor(TraverserVisitor, TypeTraverserVisitor): """Recursive traversal of both Node and Type objects.""" # Symbol nodes def visit_var(self, var: Var) -> None: self.visit_optional_type(var.type) def visit_func(self, o: FuncItem) -> None: super().visit_func(o) self.visit_optional_type(o.type) def visit_class_def(self, o: ClassDef) -> None: # TODO: Should we visit generated methods/variables as well, either here or in # TraverserVisitor? super().visit_class_def(o) info = o.info if info: for base in info.bases: base.accept(self) def visit_type_alias_expr(self, o: TypeAliasExpr) -> None: super().visit_type_alias_expr(o) o.type.accept(self) def visit_type_var_expr(self, o: TypeVarExpr) -> None: super().visit_type_var_expr(o) o.upper_bound.accept(self) for value in o.values: value.accept(self) def visit_typeddict_expr(self, o: TypedDictExpr) -> None: super().visit_typeddict_expr(o) self.visit_optional_type(o.info.typeddict_type) def visit_namedtuple_expr(self, o: NamedTupleExpr) -> None: super().visit_namedtuple_expr(o) assert o.info.tuple_type o.info.tuple_type.accept(self) def visit__promote_expr(self, o: PromoteExpr) -> None: super().visit__promote_expr(o) o.type.accept(self) def visit_newtype_expr(self, o: NewTypeExpr) -> None: super().visit_newtype_expr(o) self.visit_optional_type(o.old_type) # Statements def visit_assignment_stmt(self, o: AssignmentStmt) -> None: super().visit_assignment_stmt(o) self.visit_optional_type(o.type) def visit_for_stmt(self, o: ForStmt) -> None: super().visit_for_stmt(o) self.visit_optional_type(o.index_type) def visit_with_stmt(self, o: WithStmt) -> None: super().visit_with_stmt(o) for typ in o.analyzed_types: typ.accept(self) # Expressions def visit_cast_expr(self, o: CastExpr) -> None: super().visit_cast_expr(o) o.type.accept(self) def visit_assert_type_expr(self, o: AssertTypeExpr) -> None: super().visit_assert_type_expr(o) o.type.accept(self) def visit_type_application(self, o: TypeApplication) -> None: super().visit_type_application(o) for t in o.types: t.accept(self) # Helpers def visit_optional_type(self, t: Optional[Type]) -> None: if t: t.accept(self)