subexpr.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198
  1. """Find all subexpressions of an AST node."""
  2. from __future__ import annotations
  3. from mypy.nodes import (
  4. AssertTypeExpr,
  5. AssignmentExpr,
  6. AwaitExpr,
  7. CallExpr,
  8. CastExpr,
  9. ComparisonExpr,
  10. ConditionalExpr,
  11. DictExpr,
  12. DictionaryComprehension,
  13. Expression,
  14. GeneratorExpr,
  15. IndexExpr,
  16. LambdaExpr,
  17. ListComprehension,
  18. ListExpr,
  19. MemberExpr,
  20. Node,
  21. OpExpr,
  22. RevealExpr,
  23. SetComprehension,
  24. SetExpr,
  25. SliceExpr,
  26. StarExpr,
  27. TupleExpr,
  28. TypeApplication,
  29. UnaryExpr,
  30. YieldExpr,
  31. YieldFromExpr,
  32. )
  33. from mypy.traverser import TraverserVisitor
  34. def get_subexpressions(node: Node) -> list[Expression]:
  35. visitor = SubexpressionFinder()
  36. node.accept(visitor)
  37. return visitor.expressions
  38. class SubexpressionFinder(TraverserVisitor):
  39. def __init__(self) -> None:
  40. self.expressions: list[Expression] = []
  41. def visit_int_expr(self, o: Expression) -> None:
  42. self.add(o)
  43. def visit_name_expr(self, o: Expression) -> None:
  44. self.add(o)
  45. def visit_float_expr(self, o: Expression) -> None:
  46. self.add(o)
  47. def visit_str_expr(self, o: Expression) -> None:
  48. self.add(o)
  49. def visit_bytes_expr(self, o: Expression) -> None:
  50. self.add(o)
  51. def visit_unicode_expr(self, o: Expression) -> None:
  52. self.add(o)
  53. def visit_complex_expr(self, o: Expression) -> None:
  54. self.add(o)
  55. def visit_ellipsis(self, o: Expression) -> None:
  56. self.add(o)
  57. def visit_super_expr(self, o: Expression) -> None:
  58. self.add(o)
  59. def visit_type_var_expr(self, o: Expression) -> None:
  60. self.add(o)
  61. def visit_type_alias_expr(self, o: Expression) -> None:
  62. self.add(o)
  63. def visit_namedtuple_expr(self, o: Expression) -> None:
  64. self.add(o)
  65. def visit_typeddict_expr(self, o: Expression) -> None:
  66. self.add(o)
  67. def visit__promote_expr(self, o: Expression) -> None:
  68. self.add(o)
  69. def visit_newtype_expr(self, o: Expression) -> None:
  70. self.add(o)
  71. def visit_member_expr(self, e: MemberExpr) -> None:
  72. self.add(e)
  73. super().visit_member_expr(e)
  74. def visit_yield_from_expr(self, e: YieldFromExpr) -> None:
  75. self.add(e)
  76. super().visit_yield_from_expr(e)
  77. def visit_yield_expr(self, e: YieldExpr) -> None:
  78. self.add(e)
  79. super().visit_yield_expr(e)
  80. def visit_call_expr(self, e: CallExpr) -> None:
  81. self.add(e)
  82. super().visit_call_expr(e)
  83. def visit_op_expr(self, e: OpExpr) -> None:
  84. self.add(e)
  85. super().visit_op_expr(e)
  86. def visit_comparison_expr(self, e: ComparisonExpr) -> None:
  87. self.add(e)
  88. super().visit_comparison_expr(e)
  89. def visit_slice_expr(self, e: SliceExpr) -> None:
  90. self.add(e)
  91. super().visit_slice_expr(e)
  92. def visit_cast_expr(self, e: CastExpr) -> None:
  93. self.add(e)
  94. super().visit_cast_expr(e)
  95. def visit_assert_type_expr(self, e: AssertTypeExpr) -> None:
  96. self.add(e)
  97. super().visit_assert_type_expr(e)
  98. def visit_reveal_expr(self, e: RevealExpr) -> None:
  99. self.add(e)
  100. super().visit_reveal_expr(e)
  101. def visit_assignment_expr(self, e: AssignmentExpr) -> None:
  102. self.add(e)
  103. super().visit_assignment_expr(e)
  104. def visit_unary_expr(self, e: UnaryExpr) -> None:
  105. self.add(e)
  106. super().visit_unary_expr(e)
  107. def visit_list_expr(self, e: ListExpr) -> None:
  108. self.add(e)
  109. super().visit_list_expr(e)
  110. def visit_tuple_expr(self, e: TupleExpr) -> None:
  111. self.add(e)
  112. super().visit_tuple_expr(e)
  113. def visit_dict_expr(self, e: DictExpr) -> None:
  114. self.add(e)
  115. super().visit_dict_expr(e)
  116. def visit_set_expr(self, e: SetExpr) -> None:
  117. self.add(e)
  118. super().visit_set_expr(e)
  119. def visit_index_expr(self, e: IndexExpr) -> None:
  120. self.add(e)
  121. super().visit_index_expr(e)
  122. def visit_generator_expr(self, e: GeneratorExpr) -> None:
  123. self.add(e)
  124. super().visit_generator_expr(e)
  125. def visit_dictionary_comprehension(self, e: DictionaryComprehension) -> None:
  126. self.add(e)
  127. super().visit_dictionary_comprehension(e)
  128. def visit_list_comprehension(self, e: ListComprehension) -> None:
  129. self.add(e)
  130. super().visit_list_comprehension(e)
  131. def visit_set_comprehension(self, e: SetComprehension) -> None:
  132. self.add(e)
  133. super().visit_set_comprehension(e)
  134. def visit_conditional_expr(self, e: ConditionalExpr) -> None:
  135. self.add(e)
  136. super().visit_conditional_expr(e)
  137. def visit_type_application(self, e: TypeApplication) -> None:
  138. self.add(e)
  139. super().visit_type_application(e)
  140. def visit_lambda_expr(self, e: LambdaExpr) -> None:
  141. self.add(e)
  142. super().visit_lambda_expr(e)
  143. def visit_star_expr(self, e: StarExpr) -> None:
  144. self.add(e)
  145. super().visit_star_expr(e)
  146. def visit_await_expr(self, e: AwaitExpr) -> None:
  147. self.add(e)
  148. super().visit_await_expr(e)
  149. def add(self, e: Expression) -> None:
  150. self.expressions.append(e)