mixedtraverser.py 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  1. from __future__ import annotations
  2. from mypy.nodes import (
  3. AssertTypeExpr,
  4. AssignmentStmt,
  5. CastExpr,
  6. ClassDef,
  7. ForStmt,
  8. FuncItem,
  9. NamedTupleExpr,
  10. NewTypeExpr,
  11. PromoteExpr,
  12. TypeAliasExpr,
  13. TypeApplication,
  14. TypedDictExpr,
  15. TypeVarExpr,
  16. Var,
  17. WithStmt,
  18. )
  19. from mypy.traverser import TraverserVisitor
  20. from mypy.types import Type
  21. from mypy.typetraverser import TypeTraverserVisitor
  22. class MixedTraverserVisitor(TraverserVisitor, TypeTraverserVisitor):
  23. """Recursive traversal of both Node and Type objects."""
  24. def __init__(self) -> None:
  25. self.in_type_alias_expr = False
  26. # Symbol nodes
  27. def visit_var(self, var: Var) -> None:
  28. self.visit_optional_type(var.type)
  29. def visit_func(self, o: FuncItem) -> None:
  30. super().visit_func(o)
  31. self.visit_optional_type(o.type)
  32. def visit_class_def(self, o: ClassDef) -> None:
  33. # TODO: Should we visit generated methods/variables as well, either here or in
  34. # TraverserVisitor?
  35. super().visit_class_def(o)
  36. info = o.info
  37. if info:
  38. for base in info.bases:
  39. base.accept(self)
  40. def visit_type_alias_expr(self, o: TypeAliasExpr) -> None:
  41. super().visit_type_alias_expr(o)
  42. self.in_type_alias_expr = True
  43. o.type.accept(self)
  44. self.in_type_alias_expr = False
  45. def visit_type_var_expr(self, o: TypeVarExpr) -> None:
  46. super().visit_type_var_expr(o)
  47. o.upper_bound.accept(self)
  48. for value in o.values:
  49. value.accept(self)
  50. def visit_typeddict_expr(self, o: TypedDictExpr) -> None:
  51. super().visit_typeddict_expr(o)
  52. self.visit_optional_type(o.info.typeddict_type)
  53. def visit_namedtuple_expr(self, o: NamedTupleExpr) -> None:
  54. super().visit_namedtuple_expr(o)
  55. assert o.info.tuple_type
  56. o.info.tuple_type.accept(self)
  57. def visit__promote_expr(self, o: PromoteExpr) -> None:
  58. super().visit__promote_expr(o)
  59. o.type.accept(self)
  60. def visit_newtype_expr(self, o: NewTypeExpr) -> None:
  61. super().visit_newtype_expr(o)
  62. self.visit_optional_type(o.old_type)
  63. # Statements
  64. def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
  65. super().visit_assignment_stmt(o)
  66. self.visit_optional_type(o.type)
  67. def visit_for_stmt(self, o: ForStmt) -> None:
  68. super().visit_for_stmt(o)
  69. self.visit_optional_type(o.index_type)
  70. def visit_with_stmt(self, o: WithStmt) -> None:
  71. super().visit_with_stmt(o)
  72. for typ in o.analyzed_types:
  73. typ.accept(self)
  74. # Expressions
  75. def visit_cast_expr(self, o: CastExpr) -> None:
  76. super().visit_cast_expr(o)
  77. o.type.accept(self)
  78. def visit_assert_type_expr(self, o: AssertTypeExpr) -> None:
  79. super().visit_assert_type_expr(o)
  80. o.type.accept(self)
  81. def visit_type_application(self, o: TypeApplication) -> None:
  82. super().visit_type_application(o)
  83. for t in o.types:
  84. t.accept(self)
  85. # Helpers
  86. def visit_optional_type(self, t: Type | None) -> None:
  87. if t:
  88. t.accept(self)