traverser.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960
  1. """Generic node traverser visitor"""
  2. from __future__ import annotations
  3. from mypy_extensions import mypyc_attr
  4. from mypy.nodes import (
  5. REVEAL_TYPE,
  6. AssertStmt,
  7. AssertTypeExpr,
  8. AssignmentExpr,
  9. AssignmentStmt,
  10. AwaitExpr,
  11. Block,
  12. BreakStmt,
  13. BytesExpr,
  14. CallExpr,
  15. CastExpr,
  16. ClassDef,
  17. ComparisonExpr,
  18. ComplexExpr,
  19. ConditionalExpr,
  20. ContinueStmt,
  21. Decorator,
  22. DelStmt,
  23. DictExpr,
  24. DictionaryComprehension,
  25. EllipsisExpr,
  26. EnumCallExpr,
  27. Expression,
  28. ExpressionStmt,
  29. FloatExpr,
  30. ForStmt,
  31. FuncBase,
  32. FuncDef,
  33. FuncItem,
  34. GeneratorExpr,
  35. GlobalDecl,
  36. IfStmt,
  37. Import,
  38. ImportAll,
  39. ImportFrom,
  40. IndexExpr,
  41. IntExpr,
  42. LambdaExpr,
  43. ListComprehension,
  44. ListExpr,
  45. MatchStmt,
  46. MemberExpr,
  47. MypyFile,
  48. NamedTupleExpr,
  49. NameExpr,
  50. NewTypeExpr,
  51. Node,
  52. NonlocalDecl,
  53. OperatorAssignmentStmt,
  54. OpExpr,
  55. OverloadedFuncDef,
  56. ParamSpecExpr,
  57. PassStmt,
  58. RaiseStmt,
  59. ReturnStmt,
  60. RevealExpr,
  61. SetComprehension,
  62. SetExpr,
  63. SliceExpr,
  64. StarExpr,
  65. StrExpr,
  66. SuperExpr,
  67. TryStmt,
  68. TupleExpr,
  69. TypeAlias,
  70. TypeAliasExpr,
  71. TypeApplication,
  72. TypedDictExpr,
  73. TypeVarExpr,
  74. TypeVarTupleExpr,
  75. UnaryExpr,
  76. WhileStmt,
  77. WithStmt,
  78. YieldExpr,
  79. YieldFromExpr,
  80. )
  81. from mypy.patterns import (
  82. AsPattern,
  83. ClassPattern,
  84. MappingPattern,
  85. OrPattern,
  86. SequencePattern,
  87. SingletonPattern,
  88. StarredPattern,
  89. ValuePattern,
  90. )
  91. from mypy.visitor import NodeVisitor
  92. @mypyc_attr(allow_interpreted_subclasses=True)
  93. class TraverserVisitor(NodeVisitor[None]):
  94. """A parse tree visitor that traverses the parse tree during visiting.
  95. It does not perform any actions outside the traversal. Subclasses
  96. should override visit methods to perform actions during
  97. traversal. Calling the superclass method allows reusing the
  98. traversal implementation.
  99. """
  100. def __init__(self) -> None:
  101. pass
  102. # Visit methods
  103. def visit_mypy_file(self, o: MypyFile) -> None:
  104. for d in o.defs:
  105. d.accept(self)
  106. def visit_block(self, block: Block) -> None:
  107. for s in block.body:
  108. s.accept(self)
  109. def visit_func(self, o: FuncItem) -> None:
  110. if o.arguments is not None:
  111. for arg in o.arguments:
  112. init = arg.initializer
  113. if init is not None:
  114. init.accept(self)
  115. for arg in o.arguments:
  116. self.visit_var(arg.variable)
  117. o.body.accept(self)
  118. def visit_func_def(self, o: FuncDef) -> None:
  119. self.visit_func(o)
  120. def visit_overloaded_func_def(self, o: OverloadedFuncDef) -> None:
  121. for item in o.items:
  122. item.accept(self)
  123. if o.impl:
  124. o.impl.accept(self)
  125. def visit_class_def(self, o: ClassDef) -> None:
  126. for d in o.decorators:
  127. d.accept(self)
  128. for base in o.base_type_exprs:
  129. base.accept(self)
  130. if o.metaclass:
  131. o.metaclass.accept(self)
  132. for v in o.keywords.values():
  133. v.accept(self)
  134. o.defs.accept(self)
  135. if o.analyzed:
  136. o.analyzed.accept(self)
  137. def visit_decorator(self, o: Decorator) -> None:
  138. o.func.accept(self)
  139. o.var.accept(self)
  140. for decorator in o.decorators:
  141. decorator.accept(self)
  142. def visit_expression_stmt(self, o: ExpressionStmt) -> None:
  143. o.expr.accept(self)
  144. def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
  145. o.rvalue.accept(self)
  146. for l in o.lvalues:
  147. l.accept(self)
  148. def visit_operator_assignment_stmt(self, o: OperatorAssignmentStmt) -> None:
  149. o.rvalue.accept(self)
  150. o.lvalue.accept(self)
  151. def visit_while_stmt(self, o: WhileStmt) -> None:
  152. o.expr.accept(self)
  153. o.body.accept(self)
  154. if o.else_body:
  155. o.else_body.accept(self)
  156. def visit_for_stmt(self, o: ForStmt) -> None:
  157. o.index.accept(self)
  158. o.expr.accept(self)
  159. o.body.accept(self)
  160. if o.else_body:
  161. o.else_body.accept(self)
  162. def visit_return_stmt(self, o: ReturnStmt) -> None:
  163. if o.expr is not None:
  164. o.expr.accept(self)
  165. def visit_assert_stmt(self, o: AssertStmt) -> None:
  166. if o.expr is not None:
  167. o.expr.accept(self)
  168. if o.msg is not None:
  169. o.msg.accept(self)
  170. def visit_del_stmt(self, o: DelStmt) -> None:
  171. if o.expr is not None:
  172. o.expr.accept(self)
  173. def visit_if_stmt(self, o: IfStmt) -> None:
  174. for e in o.expr:
  175. e.accept(self)
  176. for b in o.body:
  177. b.accept(self)
  178. if o.else_body:
  179. o.else_body.accept(self)
  180. def visit_raise_stmt(self, o: RaiseStmt) -> None:
  181. if o.expr is not None:
  182. o.expr.accept(self)
  183. if o.from_expr is not None:
  184. o.from_expr.accept(self)
  185. def visit_try_stmt(self, o: TryStmt) -> None:
  186. o.body.accept(self)
  187. for i in range(len(o.types)):
  188. tp = o.types[i]
  189. if tp is not None:
  190. tp.accept(self)
  191. o.handlers[i].accept(self)
  192. for v in o.vars:
  193. if v is not None:
  194. v.accept(self)
  195. if o.else_body is not None:
  196. o.else_body.accept(self)
  197. if o.finally_body is not None:
  198. o.finally_body.accept(self)
  199. def visit_with_stmt(self, o: WithStmt) -> None:
  200. for i in range(len(o.expr)):
  201. o.expr[i].accept(self)
  202. targ = o.target[i]
  203. if targ is not None:
  204. targ.accept(self)
  205. o.body.accept(self)
  206. def visit_match_stmt(self, o: MatchStmt) -> None:
  207. o.subject.accept(self)
  208. for i in range(len(o.patterns)):
  209. o.patterns[i].accept(self)
  210. guard = o.guards[i]
  211. if guard is not None:
  212. guard.accept(self)
  213. o.bodies[i].accept(self)
  214. def visit_member_expr(self, o: MemberExpr) -> None:
  215. o.expr.accept(self)
  216. def visit_yield_from_expr(self, o: YieldFromExpr) -> None:
  217. o.expr.accept(self)
  218. def visit_yield_expr(self, o: YieldExpr) -> None:
  219. if o.expr:
  220. o.expr.accept(self)
  221. def visit_call_expr(self, o: CallExpr) -> None:
  222. o.callee.accept(self)
  223. for a in o.args:
  224. a.accept(self)
  225. if o.analyzed:
  226. o.analyzed.accept(self)
  227. def visit_op_expr(self, o: OpExpr) -> None:
  228. o.left.accept(self)
  229. o.right.accept(self)
  230. if o.analyzed is not None:
  231. o.analyzed.accept(self)
  232. def visit_comparison_expr(self, o: ComparisonExpr) -> None:
  233. for operand in o.operands:
  234. operand.accept(self)
  235. def visit_slice_expr(self, o: SliceExpr) -> None:
  236. if o.begin_index is not None:
  237. o.begin_index.accept(self)
  238. if o.end_index is not None:
  239. o.end_index.accept(self)
  240. if o.stride is not None:
  241. o.stride.accept(self)
  242. def visit_cast_expr(self, o: CastExpr) -> None:
  243. o.expr.accept(self)
  244. def visit_assert_type_expr(self, o: AssertTypeExpr) -> None:
  245. o.expr.accept(self)
  246. def visit_reveal_expr(self, o: RevealExpr) -> None:
  247. if o.kind == REVEAL_TYPE:
  248. assert o.expr is not None
  249. o.expr.accept(self)
  250. else:
  251. # RevealLocalsExpr doesn't have an inner expression
  252. pass
  253. def visit_assignment_expr(self, o: AssignmentExpr) -> None:
  254. o.target.accept(self)
  255. o.value.accept(self)
  256. def visit_unary_expr(self, o: UnaryExpr) -> None:
  257. o.expr.accept(self)
  258. def visit_list_expr(self, o: ListExpr) -> None:
  259. for item in o.items:
  260. item.accept(self)
  261. def visit_tuple_expr(self, o: TupleExpr) -> None:
  262. for item in o.items:
  263. item.accept(self)
  264. def visit_dict_expr(self, o: DictExpr) -> None:
  265. for k, v in o.items:
  266. if k is not None:
  267. k.accept(self)
  268. v.accept(self)
  269. def visit_set_expr(self, o: SetExpr) -> None:
  270. for item in o.items:
  271. item.accept(self)
  272. def visit_index_expr(self, o: IndexExpr) -> None:
  273. o.base.accept(self)
  274. o.index.accept(self)
  275. if o.analyzed:
  276. o.analyzed.accept(self)
  277. def visit_generator_expr(self, o: GeneratorExpr) -> None:
  278. for index, sequence, conditions in zip(o.indices, o.sequences, o.condlists):
  279. sequence.accept(self)
  280. index.accept(self)
  281. for cond in conditions:
  282. cond.accept(self)
  283. o.left_expr.accept(self)
  284. def visit_dictionary_comprehension(self, o: DictionaryComprehension) -> None:
  285. for index, sequence, conditions in zip(o.indices, o.sequences, o.condlists):
  286. sequence.accept(self)
  287. index.accept(self)
  288. for cond in conditions:
  289. cond.accept(self)
  290. o.key.accept(self)
  291. o.value.accept(self)
  292. def visit_list_comprehension(self, o: ListComprehension) -> None:
  293. o.generator.accept(self)
  294. def visit_set_comprehension(self, o: SetComprehension) -> None:
  295. o.generator.accept(self)
  296. def visit_conditional_expr(self, o: ConditionalExpr) -> None:
  297. o.cond.accept(self)
  298. o.if_expr.accept(self)
  299. o.else_expr.accept(self)
  300. def visit_type_application(self, o: TypeApplication) -> None:
  301. o.expr.accept(self)
  302. def visit_lambda_expr(self, o: LambdaExpr) -> None:
  303. self.visit_func(o)
  304. def visit_star_expr(self, o: StarExpr) -> None:
  305. o.expr.accept(self)
  306. def visit_await_expr(self, o: AwaitExpr) -> None:
  307. o.expr.accept(self)
  308. def visit_super_expr(self, o: SuperExpr) -> None:
  309. o.call.accept(self)
  310. def visit_as_pattern(self, o: AsPattern) -> None:
  311. if o.pattern is not None:
  312. o.pattern.accept(self)
  313. if o.name is not None:
  314. o.name.accept(self)
  315. def visit_or_pattern(self, o: OrPattern) -> None:
  316. for p in o.patterns:
  317. p.accept(self)
  318. def visit_value_pattern(self, o: ValuePattern) -> None:
  319. o.expr.accept(self)
  320. def visit_sequence_pattern(self, o: SequencePattern) -> None:
  321. for p in o.patterns:
  322. p.accept(self)
  323. def visit_starred_pattern(self, o: StarredPattern) -> None:
  324. if o.capture is not None:
  325. o.capture.accept(self)
  326. def visit_mapping_pattern(self, o: MappingPattern) -> None:
  327. for key in o.keys:
  328. key.accept(self)
  329. for value in o.values:
  330. value.accept(self)
  331. if o.rest is not None:
  332. o.rest.accept(self)
  333. def visit_class_pattern(self, o: ClassPattern) -> None:
  334. o.class_ref.accept(self)
  335. for p in o.positionals:
  336. p.accept(self)
  337. for v in o.keyword_values:
  338. v.accept(self)
  339. def visit_import(self, o: Import) -> None:
  340. for a in o.assignments:
  341. a.accept(self)
  342. def visit_import_from(self, o: ImportFrom) -> None:
  343. for a in o.assignments:
  344. a.accept(self)
  345. class ExtendedTraverserVisitor(TraverserVisitor):
  346. """This is a more flexible traverser.
  347. In addition to the base traverser it:
  348. * has visit_ methods for leaf nodes
  349. * has common method that is called for all nodes
  350. * allows to skip recursing into a node
  351. Note that this traverser still doesn't visit some internal
  352. mypy constructs like _promote expression and Var.
  353. """
  354. def visit(self, o: Node) -> bool:
  355. # If returns True, will continue to nested nodes.
  356. return True
  357. def visit_mypy_file(self, o: MypyFile) -> None:
  358. if not self.visit(o):
  359. return
  360. super().visit_mypy_file(o)
  361. # Module structure
  362. def visit_import(self, o: Import) -> None:
  363. if not self.visit(o):
  364. return
  365. super().visit_import(o)
  366. def visit_import_from(self, o: ImportFrom) -> None:
  367. if not self.visit(o):
  368. return
  369. super().visit_import_from(o)
  370. def visit_import_all(self, o: ImportAll) -> None:
  371. if not self.visit(o):
  372. return
  373. super().visit_import_all(o)
  374. # Definitions
  375. def visit_func_def(self, o: FuncDef) -> None:
  376. if not self.visit(o):
  377. return
  378. super().visit_func_def(o)
  379. def visit_overloaded_func_def(self, o: OverloadedFuncDef) -> None:
  380. if not self.visit(o):
  381. return
  382. super().visit_overloaded_func_def(o)
  383. def visit_class_def(self, o: ClassDef) -> None:
  384. if not self.visit(o):
  385. return
  386. super().visit_class_def(o)
  387. def visit_global_decl(self, o: GlobalDecl) -> None:
  388. if not self.visit(o):
  389. return
  390. super().visit_global_decl(o)
  391. def visit_nonlocal_decl(self, o: NonlocalDecl) -> None:
  392. if not self.visit(o):
  393. return
  394. super().visit_nonlocal_decl(o)
  395. def visit_decorator(self, o: Decorator) -> None:
  396. if not self.visit(o):
  397. return
  398. super().visit_decorator(o)
  399. def visit_type_alias(self, o: TypeAlias) -> None:
  400. if not self.visit(o):
  401. return
  402. super().visit_type_alias(o)
  403. # Statements
  404. def visit_block(self, block: Block) -> None:
  405. if not self.visit(block):
  406. return
  407. super().visit_block(block)
  408. def visit_expression_stmt(self, o: ExpressionStmt) -> None:
  409. if not self.visit(o):
  410. return
  411. super().visit_expression_stmt(o)
  412. def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
  413. if not self.visit(o):
  414. return
  415. super().visit_assignment_stmt(o)
  416. def visit_operator_assignment_stmt(self, o: OperatorAssignmentStmt) -> None:
  417. if not self.visit(o):
  418. return
  419. super().visit_operator_assignment_stmt(o)
  420. def visit_while_stmt(self, o: WhileStmt) -> None:
  421. if not self.visit(o):
  422. return
  423. super().visit_while_stmt(o)
  424. def visit_for_stmt(self, o: ForStmt) -> None:
  425. if not self.visit(o):
  426. return
  427. super().visit_for_stmt(o)
  428. def visit_return_stmt(self, o: ReturnStmt) -> None:
  429. if not self.visit(o):
  430. return
  431. super().visit_return_stmt(o)
  432. def visit_assert_stmt(self, o: AssertStmt) -> None:
  433. if not self.visit(o):
  434. return
  435. super().visit_assert_stmt(o)
  436. def visit_del_stmt(self, o: DelStmt) -> None:
  437. if not self.visit(o):
  438. return
  439. super().visit_del_stmt(o)
  440. def visit_if_stmt(self, o: IfStmt) -> None:
  441. if not self.visit(o):
  442. return
  443. super().visit_if_stmt(o)
  444. def visit_break_stmt(self, o: BreakStmt) -> None:
  445. if not self.visit(o):
  446. return
  447. super().visit_break_stmt(o)
  448. def visit_continue_stmt(self, o: ContinueStmt) -> None:
  449. if not self.visit(o):
  450. return
  451. super().visit_continue_stmt(o)
  452. def visit_pass_stmt(self, o: PassStmt) -> None:
  453. if not self.visit(o):
  454. return
  455. super().visit_pass_stmt(o)
  456. def visit_raise_stmt(self, o: RaiseStmt) -> None:
  457. if not self.visit(o):
  458. return
  459. super().visit_raise_stmt(o)
  460. def visit_try_stmt(self, o: TryStmt) -> None:
  461. if not self.visit(o):
  462. return
  463. super().visit_try_stmt(o)
  464. def visit_with_stmt(self, o: WithStmt) -> None:
  465. if not self.visit(o):
  466. return
  467. super().visit_with_stmt(o)
  468. def visit_match_stmt(self, o: MatchStmt) -> None:
  469. if not self.visit(o):
  470. return
  471. super().visit_match_stmt(o)
  472. # Expressions (default no-op implementation)
  473. def visit_int_expr(self, o: IntExpr) -> None:
  474. if not self.visit(o):
  475. return
  476. super().visit_int_expr(o)
  477. def visit_str_expr(self, o: StrExpr) -> None:
  478. if not self.visit(o):
  479. return
  480. super().visit_str_expr(o)
  481. def visit_bytes_expr(self, o: BytesExpr) -> None:
  482. if not self.visit(o):
  483. return
  484. super().visit_bytes_expr(o)
  485. def visit_float_expr(self, o: FloatExpr) -> None:
  486. if not self.visit(o):
  487. return
  488. super().visit_float_expr(o)
  489. def visit_complex_expr(self, o: ComplexExpr) -> None:
  490. if not self.visit(o):
  491. return
  492. super().visit_complex_expr(o)
  493. def visit_ellipsis(self, o: EllipsisExpr) -> None:
  494. if not self.visit(o):
  495. return
  496. super().visit_ellipsis(o)
  497. def visit_star_expr(self, o: StarExpr) -> None:
  498. if not self.visit(o):
  499. return
  500. super().visit_star_expr(o)
  501. def visit_name_expr(self, o: NameExpr) -> None:
  502. if not self.visit(o):
  503. return
  504. super().visit_name_expr(o)
  505. def visit_member_expr(self, o: MemberExpr) -> None:
  506. if not self.visit(o):
  507. return
  508. super().visit_member_expr(o)
  509. def visit_yield_from_expr(self, o: YieldFromExpr) -> None:
  510. if not self.visit(o):
  511. return
  512. super().visit_yield_from_expr(o)
  513. def visit_yield_expr(self, o: YieldExpr) -> None:
  514. if not self.visit(o):
  515. return
  516. super().visit_yield_expr(o)
  517. def visit_call_expr(self, o: CallExpr) -> None:
  518. if not self.visit(o):
  519. return
  520. super().visit_call_expr(o)
  521. def visit_op_expr(self, o: OpExpr) -> None:
  522. if not self.visit(o):
  523. return
  524. super().visit_op_expr(o)
  525. def visit_comparison_expr(self, o: ComparisonExpr) -> None:
  526. if not self.visit(o):
  527. return
  528. super().visit_comparison_expr(o)
  529. def visit_cast_expr(self, o: CastExpr) -> None:
  530. if not self.visit(o):
  531. return
  532. super().visit_cast_expr(o)
  533. def visit_assert_type_expr(self, o: AssertTypeExpr) -> None:
  534. if not self.visit(o):
  535. return
  536. super().visit_assert_type_expr(o)
  537. def visit_reveal_expr(self, o: RevealExpr) -> None:
  538. if not self.visit(o):
  539. return
  540. super().visit_reveal_expr(o)
  541. def visit_super_expr(self, o: SuperExpr) -> None:
  542. if not self.visit(o):
  543. return
  544. super().visit_super_expr(o)
  545. def visit_assignment_expr(self, o: AssignmentExpr) -> None:
  546. if not self.visit(o):
  547. return
  548. super().visit_assignment_expr(o)
  549. def visit_unary_expr(self, o: UnaryExpr) -> None:
  550. if not self.visit(o):
  551. return
  552. super().visit_unary_expr(o)
  553. def visit_list_expr(self, o: ListExpr) -> None:
  554. if not self.visit(o):
  555. return
  556. super().visit_list_expr(o)
  557. def visit_dict_expr(self, o: DictExpr) -> None:
  558. if not self.visit(o):
  559. return
  560. super().visit_dict_expr(o)
  561. def visit_tuple_expr(self, o: TupleExpr) -> None:
  562. if not self.visit(o):
  563. return
  564. super().visit_tuple_expr(o)
  565. def visit_set_expr(self, o: SetExpr) -> None:
  566. if not self.visit(o):
  567. return
  568. super().visit_set_expr(o)
  569. def visit_index_expr(self, o: IndexExpr) -> None:
  570. if not self.visit(o):
  571. return
  572. super().visit_index_expr(o)
  573. def visit_type_application(self, o: TypeApplication) -> None:
  574. if not self.visit(o):
  575. return
  576. super().visit_type_application(o)
  577. def visit_lambda_expr(self, o: LambdaExpr) -> None:
  578. if not self.visit(o):
  579. return
  580. super().visit_lambda_expr(o)
  581. def visit_list_comprehension(self, o: ListComprehension) -> None:
  582. if not self.visit(o):
  583. return
  584. super().visit_list_comprehension(o)
  585. def visit_set_comprehension(self, o: SetComprehension) -> None:
  586. if not self.visit(o):
  587. return
  588. super().visit_set_comprehension(o)
  589. def visit_dictionary_comprehension(self, o: DictionaryComprehension) -> None:
  590. if not self.visit(o):
  591. return
  592. super().visit_dictionary_comprehension(o)
  593. def visit_generator_expr(self, o: GeneratorExpr) -> None:
  594. if not self.visit(o):
  595. return
  596. super().visit_generator_expr(o)
  597. def visit_slice_expr(self, o: SliceExpr) -> None:
  598. if not self.visit(o):
  599. return
  600. super().visit_slice_expr(o)
  601. def visit_conditional_expr(self, o: ConditionalExpr) -> None:
  602. if not self.visit(o):
  603. return
  604. super().visit_conditional_expr(o)
  605. def visit_type_var_expr(self, o: TypeVarExpr) -> None:
  606. if not self.visit(o):
  607. return
  608. super().visit_type_var_expr(o)
  609. def visit_paramspec_expr(self, o: ParamSpecExpr) -> None:
  610. if not self.visit(o):
  611. return
  612. super().visit_paramspec_expr(o)
  613. def visit_type_var_tuple_expr(self, o: TypeVarTupleExpr) -> None:
  614. if not self.visit(o):
  615. return
  616. super().visit_type_var_tuple_expr(o)
  617. def visit_type_alias_expr(self, o: TypeAliasExpr) -> None:
  618. if not self.visit(o):
  619. return
  620. super().visit_type_alias_expr(o)
  621. def visit_namedtuple_expr(self, o: NamedTupleExpr) -> None:
  622. if not self.visit(o):
  623. return
  624. super().visit_namedtuple_expr(o)
  625. def visit_enum_call_expr(self, o: EnumCallExpr) -> None:
  626. if not self.visit(o):
  627. return
  628. super().visit_enum_call_expr(o)
  629. def visit_typeddict_expr(self, o: TypedDictExpr) -> None:
  630. if not self.visit(o):
  631. return
  632. super().visit_typeddict_expr(o)
  633. def visit_newtype_expr(self, o: NewTypeExpr) -> None:
  634. if not self.visit(o):
  635. return
  636. super().visit_newtype_expr(o)
  637. def visit_await_expr(self, o: AwaitExpr) -> None:
  638. if not self.visit(o):
  639. return
  640. super().visit_await_expr(o)
  641. # Patterns
  642. def visit_as_pattern(self, o: AsPattern) -> None:
  643. if not self.visit(o):
  644. return
  645. super().visit_as_pattern(o)
  646. def visit_or_pattern(self, o: OrPattern) -> None:
  647. if not self.visit(o):
  648. return
  649. super().visit_or_pattern(o)
  650. def visit_value_pattern(self, o: ValuePattern) -> None:
  651. if not self.visit(o):
  652. return
  653. super().visit_value_pattern(o)
  654. def visit_singleton_pattern(self, o: SingletonPattern) -> None:
  655. if not self.visit(o):
  656. return
  657. super().visit_singleton_pattern(o)
  658. def visit_sequence_pattern(self, o: SequencePattern) -> None:
  659. if not self.visit(o):
  660. return
  661. super().visit_sequence_pattern(o)
  662. def visit_starred_pattern(self, o: StarredPattern) -> None:
  663. if not self.visit(o):
  664. return
  665. super().visit_starred_pattern(o)
  666. def visit_mapping_pattern(self, o: MappingPattern) -> None:
  667. if not self.visit(o):
  668. return
  669. super().visit_mapping_pattern(o)
  670. def visit_class_pattern(self, o: ClassPattern) -> None:
  671. if not self.visit(o):
  672. return
  673. super().visit_class_pattern(o)
  674. class ReturnSeeker(TraverserVisitor):
  675. def __init__(self) -> None:
  676. self.found = False
  677. def visit_return_stmt(self, o: ReturnStmt) -> None:
  678. if o.expr is None or isinstance(o.expr, NameExpr) and o.expr.name == "None":
  679. return
  680. self.found = True
  681. def has_return_statement(fdef: FuncBase) -> bool:
  682. """Find if a function has a non-trivial return statement.
  683. Plain 'return' and 'return None' don't count.
  684. """
  685. seeker = ReturnSeeker()
  686. fdef.accept(seeker)
  687. return seeker.found
  688. class FuncCollectorBase(TraverserVisitor):
  689. def __init__(self) -> None:
  690. self.inside_func = False
  691. def visit_func_def(self, defn: FuncDef) -> None:
  692. if not self.inside_func:
  693. self.inside_func = True
  694. super().visit_func_def(defn)
  695. self.inside_func = False
  696. class YieldSeeker(FuncCollectorBase):
  697. def __init__(self) -> None:
  698. super().__init__()
  699. self.found = False
  700. def visit_yield_expr(self, o: YieldExpr) -> None:
  701. self.found = True
  702. def has_yield_expression(fdef: FuncBase) -> bool:
  703. seeker = YieldSeeker()
  704. fdef.accept(seeker)
  705. return seeker.found
  706. class YieldFromSeeker(FuncCollectorBase):
  707. def __init__(self) -> None:
  708. super().__init__()
  709. self.found = False
  710. def visit_yield_from_expr(self, o: YieldFromExpr) -> None:
  711. self.found = True
  712. def has_yield_from_expression(fdef: FuncBase) -> bool:
  713. seeker = YieldFromSeeker()
  714. fdef.accept(seeker)
  715. return seeker.found
  716. class AwaitSeeker(TraverserVisitor):
  717. def __init__(self) -> None:
  718. super().__init__()
  719. self.found = False
  720. def visit_await_expr(self, o: AwaitExpr) -> None:
  721. self.found = True
  722. def has_await_expression(expr: Expression) -> bool:
  723. seeker = AwaitSeeker()
  724. expr.accept(seeker)
  725. return seeker.found
  726. class ReturnCollector(FuncCollectorBase):
  727. def __init__(self) -> None:
  728. super().__init__()
  729. self.return_statements: list[ReturnStmt] = []
  730. def visit_return_stmt(self, stmt: ReturnStmt) -> None:
  731. self.return_statements.append(stmt)
  732. def all_return_statements(node: Node) -> list[ReturnStmt]:
  733. v = ReturnCollector()
  734. node.accept(v)
  735. return v.return_statements
  736. class YieldCollector(FuncCollectorBase):
  737. def __init__(self) -> None:
  738. super().__init__()
  739. self.in_assignment = False
  740. self.yield_expressions: list[tuple[YieldExpr, bool]] = []
  741. def visit_assignment_stmt(self, stmt: AssignmentStmt) -> None:
  742. self.in_assignment = True
  743. super().visit_assignment_stmt(stmt)
  744. self.in_assignment = False
  745. def visit_yield_expr(self, expr: YieldExpr) -> None:
  746. self.yield_expressions.append((expr, self.in_assignment))
  747. def all_yield_expressions(node: Node) -> list[tuple[YieldExpr, bool]]:
  748. v = YieldCollector()
  749. node.accept(v)
  750. return v.yield_expressions
  751. class YieldFromCollector(FuncCollectorBase):
  752. def __init__(self) -> None:
  753. super().__init__()
  754. self.in_assignment = False
  755. self.yield_from_expressions: list[tuple[YieldFromExpr, bool]] = []
  756. def visit_assignment_stmt(self, stmt: AssignmentStmt) -> None:
  757. self.in_assignment = True
  758. super().visit_assignment_stmt(stmt)
  759. self.in_assignment = False
  760. def visit_yield_from_expr(self, expr: YieldFromExpr) -> None:
  761. self.yield_from_expressions.append((expr, self.in_assignment))
  762. def all_yield_from_expressions(node: Node) -> list[tuple[YieldFromExpr, bool]]:
  763. v = YieldFromCollector()
  764. node.accept(v)
  765. return v.yield_from_expressions