Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2#
3# This source code is licensed under the MIT license found in the
4# LICENSE file in the root directory of this source tree.
6from inspect import ismethod, signature
7from typing import (
8 Any,
9 Callable,
10 cast,
11 Dict,
12 get_type_hints,
13 List,
14 Optional,
15 Sequence,
16 Set,
17 Tuple,
18 Type,
19 Union,
20)
22import libcst as cst
23from libcst import CSTTransformer, CSTVisitor
24from libcst._types import CSTNodeT
25from libcst.matchers._decorators import (
26 CONSTRUCTED_LEAVE_MATCHER_ATTR,
27 CONSTRUCTED_VISIT_MATCHER_ATTR,
28 VISIT_NEGATIVE_MATCHER_ATTR,
29 VISIT_POSITIVE_MATCHER_ATTR,
30)
31from libcst.matchers._matcher_base import (
32 AllOf,
33 AtLeastN,
34 AtMostN,
35 BaseMatcherNode,
36 extract,
37 extractall,
38 findall,
39 matches,
40 MatchIfTrue,
41 MatchMetadata,
42 MatchMetadataIfTrue,
43 OneOf,
44 replace,
45)
46from libcst.matchers._return_types import TYPED_FUNCTION_RETURN_MAPPING
48try:
49 # PEP 604 unions, in Python 3.10+
50 from types import UnionType
51except ImportError:
52 # We use this for isinstance; no annotation will be an instance of this
53 class UnionType:
54 pass
57CONCRETE_METHODS: Set[str] = {
58 *{f"visit_{cls.__name__}" for cls in TYPED_FUNCTION_RETURN_MAPPING},
59 *{f"leave_{cls.__name__}" for cls in TYPED_FUNCTION_RETURN_MAPPING},
60}
63def is_property(obj: object, attr_name: str) -> bool:
64 """Check if obj.attr is a property without evaluating it."""
65 return isinstance(getattr(type(obj), attr_name, None), property)
68# pyre-ignore We don't care about Any here, its not exposed.
69def _match_decorator_unpickler(kwargs: Any) -> "MatchDecoratorMismatch":
70 return MatchDecoratorMismatch(**kwargs)
73class MatchDecoratorMismatch(Exception):
74 def __init__(self, func: str, message: str) -> None:
75 super().__init__(f"Invalid function signature for {func}: {message}")
76 self.func = func
77 self.message = message
79 def __reduce__(
80 self,
81 ) -> Tuple[Callable[..., "MatchDecoratorMismatch"], Tuple[object, ...]]:
82 return (
83 _match_decorator_unpickler,
84 ({"func": self.func, "message": self.message},),
85 )
88def _get_possible_match_classes(matcher: BaseMatcherNode) -> List[Type[cst.CSTNode]]:
89 if isinstance(matcher, (OneOf, AllOf)):
90 return [getattr(cst, m.__class__.__name__) for m in matcher.options]
91 else:
92 return [getattr(cst, matcher.__class__.__name__)]
95def _annotation_is_union(annotation: object) -> bool:
96 return (
97 isinstance(annotation, UnionType)
98 or getattr(annotation, "__origin__", None) is Union
99 )
102def _get_possible_annotated_classes(annotation: object) -> List[Type[object]]:
103 if _annotation_is_union(annotation):
104 return getattr(annotation, "__args__", [])
105 else:
106 return [cast(Type[object], annotation)]
109def _get_valid_leave_annotations_for_classes(
110 classes: Sequence[Type[cst.CSTNode]],
111) -> Set[Type[object]]:
112 retval: Set[Type[object]] = set()
114 for cls in classes:
115 # Look up the leave annotation for each class, combine them so we get a list of
116 # all possible valid return annotations. Its not really possible for us (or
117 # pyre) to fully enforce return types given the presence of OneOf/AllOf matchers, so
118 # we do the best we can by taking a union of all valid return annotations.
119 retval.update(
120 _get_possible_annotated_classes(TYPED_FUNCTION_RETURN_MAPPING[cls])
121 )
123 return retval
126def _verify_return_annotation(
127 possible_match_classes: Sequence[Type[cst.CSTNode]],
128 # pyre-ignore We only care that meth is callable.
129 meth: Callable[..., Any],
130 decorator_name: str,
131 *,
132 expected_none: bool,
133) -> None:
134 type_hints = get_type_hints(meth)
135 if expected_none:
136 # Simply look for any annotation at all and if it exists, verify that
137 # it is "None".
138 if type_hints.get("return", type(None)) is not type(None): # noqa: E721
139 raise MatchDecoratorMismatch(
140 meth.__qualname__,
141 f"@{decorator_name} should only decorate functions that do "
142 + "not return.",
143 )
144 else:
145 if "return" not in type_hints:
146 # Can't check this, type annotation not supplied.
147 return
149 possible_annotated_classes = _get_possible_annotated_classes(
150 type_hints["return"]
151 )
152 possible_returns = _get_valid_leave_annotations_for_classes(
153 possible_match_classes
154 )
156 # Look at the union of specified return annotation, make sure that
157 # they are all subclasses of the original leave_<Node> return
158 # annotations. This catches when somebody tries to return a new node
159 # that we know can't fit where the existing node was in the tree.
160 for ret in possible_annotated_classes:
161 for annotation in possible_returns:
162 if issubclass(ret, annotation):
163 # This annotation is a superclass of the possible match,
164 # so we know that the types are correct.
165 break
166 else:
167 # The current ret was not a subclass of any of the annotated
168 # return types.
169 raise MatchDecoratorMismatch(
170 meth.__qualname__,
171 f"@{decorator_name} decorated function cannot return "
172 + f"the type {ret.__name__}.",
173 )
176def _verify_parameter_annotations(
177 possible_match_classes: Sequence[Type[cst.CSTNode]],
178 # pyre-ignore We only care that meth is callable.
179 meth: Callable[..., Any],
180 decorator_name: str,
181 *,
182 expected_param_count: int,
183) -> None:
184 # First, verify that the number of parameters is sane.
185 meth_signature = signature(meth)
186 if len(meth_signature.parameters) != expected_param_count:
187 raise MatchDecoratorMismatch(
188 meth.__qualname__,
189 f"@{decorator_name} should decorate functions which take "
190 + f"{expected_param_count} parameter"
191 + ("s" if expected_param_count > 1 else ""),
192 )
194 # Finally, for each parameter, make sure that the annotation includes
195 # each of the classes that might appear given the match string. This
196 # can be done in the simple case by just specifying the correct cst node
197 # type. For complex matches that use OneOf/AllOf, this could be a base class
198 # that encompases all possible matches, or a union.
199 params = [v for k, v in get_type_hints(meth).items() if k != "return"]
200 for param in params:
201 # Go through each possible matcher, and make sure that the annotation
202 # for types is a superclass of each matcher.
203 possible_annotated_classes = _get_possible_annotated_classes(param)
204 for match in possible_match_classes:
205 for annotation in possible_annotated_classes:
206 if issubclass(match, annotation):
207 # This annotation is a superclass of the possible match,
208 # so we know that the types are correct.
209 break
210 else:
211 # The current match was not a subclass of any of the annotated
212 # types.
213 raise MatchDecoratorMismatch(
214 meth.__qualname__,
215 f"@{decorator_name} can be called with {match.__name__} "
216 + "but the decorated function parameter annotations do "
217 + "not include this type.",
218 )
221def _check_types(
222 # pyre-ignore We don't care about the type of sequence, just that its callable.
223 decoratormap: Dict[BaseMatcherNode, Sequence[Callable[..., Any]]],
224 decorator_name: str,
225 *,
226 expected_param_count: int,
227 expected_none_return: bool,
228) -> None:
229 for matcher, methods in decoratormap.items():
230 # Given the matcher class we have, get the list of possible cst nodes that
231 # could be passed to the functionis we wrap.
232 possible_match_classes = _get_possible_match_classes(matcher)
233 has_invalid_top_level = any(
234 isinstance(m, (AtLeastN, AtMostN, MatchIfTrue))
235 for m in possible_match_classes
236 )
238 # Now, loop through each function we wrap and verify that the type signature
239 # is valid.
240 for meth in methods:
241 # First thing first, make sure this isn't wrapping an inner class.
242 if not ismethod(meth):
243 raise MatchDecoratorMismatch(
244 meth.__qualname__,
245 "Matcher decorators should only be used on methods of "
246 + "MatcherDecoratableTransformer or "
247 + "MatcherDecoratableVisitor",
248 )
249 if has_invalid_top_level:
250 raise MatchDecoratorMismatch(
251 meth.__qualname__,
252 "The root matcher in a matcher decorator cannot be an "
253 + "AtLeastN, AtMostN or MatchIfTrue matcher",
254 )
256 # Now, check that the return annotation is valid.
257 _verify_return_annotation(
258 possible_match_classes,
259 meth,
260 decorator_name,
261 expected_none=expected_none_return,
262 )
264 # Finally, check that the parameter annotations are valid.
265 _verify_parameter_annotations(
266 possible_match_classes,
267 meth,
268 decorator_name,
269 expected_param_count=expected_param_count,
270 )
273def _gather_matchers(obj: object) -> Dict[BaseMatcherNode, Optional[cst.CSTNode]]:
274 """
275 Set of gating matchers that we need to track and evaluate. We use these
276 in conjunction with the call_if_inside and call_if_not_inside decorators
277 to determine whether to call a visit/leave function.
278 """
280 visit_matchers: Dict[BaseMatcherNode, Optional[cst.CSTNode]] = {}
282 for attr_name in dir(obj):
283 if not is_property(obj, attr_name):
284 func = getattr(obj, attr_name)
285 for matcher in getattr(func, VISIT_POSITIVE_MATCHER_ATTR, []):
286 visit_matchers[cast(BaseMatcherNode, matcher)] = None
287 for matcher in getattr(func, VISIT_NEGATIVE_MATCHER_ATTR, []):
288 visit_matchers[cast(BaseMatcherNode, matcher)] = None
290 return visit_matchers
293def _assert_not_concrete(
294 decorator_name: str, func: Callable[[cst.CSTNode], None]
295) -> None:
296 if func.__name__ in CONCRETE_METHODS:
297 raise MatchDecoratorMismatch(
298 func.__qualname__,
299 f"@{decorator_name} should not decorate functions that are concrete "
300 + "visit or leave methods.",
301 )
304def _gather_constructed_visit_funcs(
305 obj: object,
306) -> Dict[BaseMatcherNode, Sequence[Callable[[cst.CSTNode], None]]]:
307 constructed_visitors: Dict[
308 BaseMatcherNode, Sequence[Callable[[cst.CSTNode], None]]
309 ] = {}
311 for funcname in dir(obj):
312 if is_property(obj, funcname):
313 continue
314 possible_func = getattr(obj, funcname)
315 if not ismethod(possible_func):
316 continue
317 func = cast(Callable[[cst.CSTNode], None], possible_func)
318 matchers = getattr(func, CONSTRUCTED_VISIT_MATCHER_ATTR, [])
319 if matchers:
320 # Make sure that we aren't accidentally putting a @visit on a visit_Node.
321 _assert_not_concrete("visit", func)
322 for matcher in matchers:
323 casted_matcher = cast(BaseMatcherNode, matcher)
324 constructed_visitors[casted_matcher] = (
325 *constructed_visitors.get(casted_matcher, ()),
326 func,
327 )
329 return constructed_visitors
332# pyre-ignore: There is no reasonable way to type this, so ignore the Any type. This
333# is because the leave_* methods have a different signature depending on whether they
334# are in a MatcherDecoratableTransformer or a MatcherDecoratableVisitor.
335def _gather_constructed_leave_funcs(
336 obj: object,
337) -> Dict[BaseMatcherNode, Sequence[Callable[..., Any]]]:
338 constructed_visitors: Dict[
339 BaseMatcherNode, Sequence[Callable[[cst.CSTNode], None]]
340 ] = {}
342 for funcname in dir(obj):
343 if is_property(obj, funcname):
344 continue
345 possible_func = getattr(obj, funcname)
346 if not ismethod(possible_func):
347 continue
348 func = cast(Callable[[cst.CSTNode], None], possible_func)
349 matchers = getattr(func, CONSTRUCTED_LEAVE_MATCHER_ATTR, [])
350 if matchers:
351 # Make sure that we aren't accidentally putting a @leave on a leave_Node.
352 _assert_not_concrete("leave", func)
353 for matcher in matchers:
354 casted_matcher = cast(BaseMatcherNode, matcher)
355 constructed_visitors[casted_matcher] = (
356 *constructed_visitors.get(casted_matcher, ()),
357 func,
358 )
360 return constructed_visitors
363def _visit_matchers(
364 matchers: Dict[BaseMatcherNode, Optional[cst.CSTNode]],
365 node: cst.CSTNode,
366 metadata_resolver: cst.MetadataDependent,
367) -> Dict[BaseMatcherNode, Optional[cst.CSTNode]]:
368 new_matchers: Dict[BaseMatcherNode, Optional[cst.CSTNode]] = {}
369 for matcher, existing_node in matchers.items():
370 # We don't care about visiting matchers that are already true.
371 if existing_node is None and matches(
372 node, matcher, metadata_resolver=metadata_resolver
373 ):
374 # This node matches! Remember which node it was so we can
375 # cancel it later.
376 new_matchers[matcher] = node
377 else:
378 new_matchers[matcher] = existing_node
379 return new_matchers
382def _leave_matchers(
383 matchers: Dict[BaseMatcherNode, Optional[cst.CSTNode]], node: cst.CSTNode
384) -> Dict[BaseMatcherNode, Optional[cst.CSTNode]]:
385 new_matchers: Dict[BaseMatcherNode, Optional[cst.CSTNode]] = {}
386 for matcher, existing_node in matchers.items():
387 if node is existing_node:
388 # This node matches, so we are no longer inside it.
389 new_matchers[matcher] = None
390 else:
391 # We aren't leaving this node.
392 new_matchers[matcher] = existing_node
393 return new_matchers
396def _all_positive_matchers_true(
397 all_matchers: Dict[BaseMatcherNode, Optional[cst.CSTNode]], obj: object
398) -> bool:
399 requested_matchers = getattr(obj, VISIT_POSITIVE_MATCHER_ATTR, [])
400 for matcher in requested_matchers:
401 if all_matchers[matcher] is None:
402 # The passed in object has been decorated with a matcher that isn't
403 # active.
404 return False
405 return True
408def _all_negative_matchers_false(
409 all_matchers: Dict[BaseMatcherNode, Optional[cst.CSTNode]], obj: object
410) -> bool:
411 requested_matchers = getattr(obj, VISIT_NEGATIVE_MATCHER_ATTR, [])
412 for matcher in requested_matchers:
413 if all_matchers[matcher] is not None:
414 # The passed in object has been decorated with a matcher that is active.
415 return False
416 return True
419def _should_allow_visit(
420 all_matchers: Dict[BaseMatcherNode, Optional[cst.CSTNode]], obj: object
421) -> bool:
422 return _all_positive_matchers_true(
423 all_matchers, obj
424 ) and _all_negative_matchers_false(all_matchers, obj)
427def _visit_constructed_funcs(
428 visit_funcs: Dict[BaseMatcherNode, Sequence[Callable[[cst.CSTNode], None]]],
429 all_matchers: Dict[BaseMatcherNode, Optional[cst.CSTNode]],
430 node: cst.CSTNode,
431 metadata_resolver: cst.MetadataDependent,
432) -> None:
433 for matcher, visit_funcs in visit_funcs.items():
434 if matches(node, matcher, metadata_resolver=metadata_resolver):
435 for visit_func in visit_funcs:
436 if _should_allow_visit(all_matchers, visit_func):
437 visit_func(node)
440class MatcherDecoratableTransformer(CSTTransformer):
441 """
442 This class provides all of the features of a :class:`libcst.CSTTransformer`, and
443 additionally supports various decorators to control when methods get called when
444 traversing a tree. Use this instead of a :class:`libcst.CSTTransformer` if you
445 wish to do more powerful decorator-based visiting.
446 """
448 def __init__(self) -> None:
449 CSTTransformer.__init__(self)
450 self.__matchers: Optional[Dict[BaseMatcherNode, Optional[cst.CSTNode]]] = None
451 # Mapping of matchers to functions. If in the course of visiting the tree,
452 # a node matches one of these matchers, the corresponding function will be
453 # called as if it was a visit_* method.
454 self._extra_visit_funcs: Dict[
455 BaseMatcherNode, Sequence[Callable[[cst.CSTNode], None]]
456 ] = _gather_constructed_visit_funcs(self)
457 # Mapping of matchers to functions. If in the course of leaving the tree,
458 # a node matches one of these matchers, the corresponding function will be
459 # called as if it was a leave_* method.
460 self._extra_leave_funcs: Dict[
461 BaseMatcherNode,
462 Sequence[
463 Callable[
464 [cst.CSTNode, cst.CSTNode], Union[cst.CSTNode, cst.RemovalSentinel]
465 ]
466 ],
467 ] = _gather_constructed_leave_funcs(self)
468 # Make sure visit/leave functions constructed with @visit and @leave decorators
469 # have correct type annotations.
470 _check_types(
471 self._extra_visit_funcs,
472 "visit",
473 expected_param_count=1,
474 expected_none_return=True,
475 )
476 _check_types(
477 self._extra_leave_funcs,
478 "leave",
479 expected_param_count=2,
480 expected_none_return=False,
481 )
483 @property
484 def _matchers(self) -> Dict[BaseMatcherNode, Optional[cst.CSTNode]]:
485 if self.__matchers is None:
486 self.__matchers = _gather_matchers(self)
487 return self.__matchers
489 @_matchers.setter
490 def _matchers(self, value: Dict[BaseMatcherNode, Optional[cst.CSTNode]]) -> None:
491 self.__matchers = value
493 def on_visit(self, node: cst.CSTNode) -> bool:
494 # First, evaluate any matchers that we have which we are not inside already.
495 self._matchers = _visit_matchers(self._matchers, node, self)
497 # Now, call any visitors that were hooked using a visit decorator.
498 _visit_constructed_funcs(self._extra_visit_funcs, self._matchers, node, self)
500 # Now, evaluate whether this current function has any matchers it requires.
501 if not _should_allow_visit(
502 self._matchers, getattr(self, f"visit_{type(node).__name__}", None)
503 ):
504 # We shouldn't visit this directly. However, we should continue
505 # visiting its children.
506 return True
508 # Either the visit_func doesn't exist, we have no matchers, or we passed all
509 # matchers. In either case, just call the superclass behavior.
510 return CSTTransformer.on_visit(self, node)
512 def on_leave(
513 self, original_node: CSTNodeT, updated_node: CSTNodeT
514 ) -> Union[CSTNodeT, cst.RemovalSentinel]:
515 # First, evaluate whether this current function has a decorator on it.
516 if _should_allow_visit(
517 self._matchers, getattr(self, f"leave_{type(original_node).__name__}", None)
518 ):
519 retval = CSTTransformer.on_leave(self, original_node, updated_node)
520 else:
521 retval = updated_node
523 # Now, call any visitors that were hooked using a leave decorator.
524 for matcher, leave_funcs in reversed(list(self._extra_leave_funcs.items())):
525 if not self.matches(original_node, matcher):
526 continue
527 for leave_func in leave_funcs:
528 if _should_allow_visit(self._matchers, leave_func) and isinstance(
529 retval, cst.CSTNode
530 ):
531 retval = leave_func(original_node, retval)
533 # Now, see if we have any matchers we should deactivate.
534 self._matchers = _leave_matchers(self._matchers, original_node)
536 # pyre-ignore The return value of on_leave is subtly wrong in that we can
537 # actually return any value that passes this node's parent's constructor
538 # validation. Fixing this is beyond the scope of this file, and would involve
539 # forcing a lot of ensure_type() checks across the codebase.
540 return retval
542 def on_visit_attribute(self, node: cst.CSTNode, attribute: str) -> None:
543 # Evaluate whether this current function has a decorator on it.
544 if _should_allow_visit(
545 self._matchers,
546 getattr(self, f"visit_{type(node).__name__}_{attribute}", None),
547 ):
548 # Either the visit_func doesn't exist, we have no matchers, or we passed all
549 # matchers. In either case, just call the superclass behavior.
550 return CSTTransformer.on_visit_attribute(self, node, attribute)
552 def on_leave_attribute(self, original_node: cst.CSTNode, attribute: str) -> None:
553 # Evaluate whether this current function has a decorator on it.
554 if _should_allow_visit(
555 self._matchers,
556 getattr(self, f"leave_{type(original_node).__name__}_{attribute}", None),
557 ):
558 # Either the visit_func doesn't exist, we have no matchers, or we passed all
559 # matchers. In either case, just call the superclass behavior.
560 CSTTransformer.on_leave_attribute(self, original_node, attribute)
562 def matches(
563 self,
564 node: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
565 matcher: BaseMatcherNode,
566 ) -> bool:
567 """
568 A convenience method to call :func:`~libcst.matchers.matches` without requiring
569 an explicit parameter for metadata. Since our instance is an instance of
570 :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see
571 documentation for :func:`~libcst.matchers.matches` as it is identical to this
572 function.
573 """
574 return matches(node, matcher, metadata_resolver=self)
576 def findall(
577 self,
578 tree: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
579 matcher: Union[
580 BaseMatcherNode,
581 MatchIfTrue[cst.CSTNode],
582 MatchMetadata,
583 MatchMetadataIfTrue,
584 ],
585 ) -> Sequence[cst.CSTNode]:
586 """
587 A convenience method to call :func:`~libcst.matchers.findall` without requiring
588 an explicit parameter for metadata. Since our instance is an instance of
589 :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see
590 documentation for :func:`~libcst.matchers.findall` as it is identical to this
591 function.
592 """
593 return findall(tree, matcher, metadata_resolver=self)
595 def extract(
596 self,
597 node: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
598 matcher: BaseMatcherNode,
599 ) -> Optional[Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]]]:
600 """
601 A convenience method to call :func:`~libcst.matchers.extract` without requiring
602 an explicit parameter for metadata. Since our instance is an instance of
603 :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see
604 documentation for :func:`~libcst.matchers.extract` as it is identical to this
605 function.
606 """
607 return extract(node, matcher, metadata_resolver=self)
609 def extractall(
610 self,
611 tree: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
612 matcher: Union[
613 BaseMatcherNode,
614 MatchIfTrue[cst.CSTNode],
615 MatchMetadata,
616 MatchMetadataIfTrue,
617 ],
618 ) -> Sequence[Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]]]:
619 """
620 A convenience method to call :func:`~libcst.matchers.extractall` without requiring
621 an explicit parameter for metadata. Since our instance is an instance of
622 :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see
623 documentation for :func:`~libcst.matchers.extractall` as it is identical to this
624 function.
625 """
626 return extractall(tree, matcher, metadata_resolver=self)
628 def replace(
629 self,
630 tree: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
631 matcher: Union[
632 BaseMatcherNode,
633 MatchIfTrue[cst.CSTNode],
634 MatchMetadata,
635 MatchMetadataIfTrue,
636 ],
637 replacement: Union[
638 cst.MaybeSentinel,
639 cst.RemovalSentinel,
640 cst.CSTNode,
641 Callable[
642 [cst.CSTNode, Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]]],
643 Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
644 ],
645 ],
646 ) -> Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode]:
647 """
648 A convenience method to call :func:`~libcst.matchers.replace` without requiring
649 an explicit parameter for metadata. Since our instance is an instance of
650 :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see
651 documentation for :func:`~libcst.matchers.replace` as it is identical to this
652 function.
653 """
654 return replace(tree, matcher, replacement, metadata_resolver=self)
657class MatcherDecoratableVisitor(CSTVisitor):
658 """
659 This class provides all of the features of a :class:`libcst.CSTVisitor`, and
660 additionally supports various decorators to control when methods get called
661 when traversing a tree. Use this instead of a :class:`libcst.CSTVisitor` if
662 you wish to do more powerful decorator-based visiting.
663 """
665 def __init__(self) -> None:
666 CSTVisitor.__init__(self)
667 self.__matchers: Optional[Dict[BaseMatcherNode, Optional[cst.CSTNode]]] = None
668 # Mapping of matchers to functions. If in the course of visiting the tree,
669 # a node matches one of these matchers, the corresponding function will be
670 # called as if it was a visit_* method.
671 self._extra_visit_funcs: Dict[
672 BaseMatcherNode, Sequence[Callable[[cst.CSTNode], None]]
673 ] = _gather_constructed_visit_funcs(self)
674 # Mapping of matchers to functions. If in the course of leaving the tree,
675 # a node matches one of these matchers, the corresponding function will be
676 # called as if it was a leave_* method.
677 self._extra_leave_funcs: Dict[
678 BaseMatcherNode, Sequence[Callable[[cst.CSTNode], None]]
679 ] = _gather_constructed_leave_funcs(self)
680 # Make sure visit/leave functions constructed with @visit and @leave decorators
681 # have correct type annotations.
682 _check_types(
683 self._extra_visit_funcs,
684 "visit",
685 expected_param_count=1,
686 expected_none_return=True,
687 )
688 _check_types(
689 self._extra_leave_funcs,
690 "leave",
691 expected_param_count=1,
692 expected_none_return=True,
693 )
695 @property
696 def _matchers(self) -> Dict[BaseMatcherNode, Optional[cst.CSTNode]]:
697 if self.__matchers is None:
698 self.__matchers = _gather_matchers(self)
699 return self.__matchers
701 @_matchers.setter
702 def _matchers(self, value: Dict[BaseMatcherNode, Optional[cst.CSTNode]]) -> None:
703 self.__matchers = value
705 def on_visit(self, node: cst.CSTNode) -> bool:
706 # First, evaluate any matchers that we have which we are not inside already.
707 self._matchers = _visit_matchers(self._matchers, node, self)
709 # Now, call any visitors that were hooked using a visit decorator.
710 _visit_constructed_funcs(self._extra_visit_funcs, self._matchers, node, self)
712 # Now, evaluate whether this current function has a decorator on it.
713 if not _should_allow_visit(
714 self._matchers, getattr(self, f"visit_{type(node).__name__}", None)
715 ):
716 # We shouldn't visit this directly. However, we should continue
717 # visiting its children.
718 return True
720 # Either the visit_func doesn't exist, we have no matchers, or we passed all
721 # matchers. In either case, just call the superclass behavior.
722 return CSTVisitor.on_visit(self, node)
724 def on_leave(self, original_node: cst.CSTNode) -> None:
725 # First, evaluate whether this current function has a decorator on it.
726 if _should_allow_visit(
727 self._matchers, getattr(self, f"leave_{type(original_node).__name__}", None)
728 ):
729 CSTVisitor.on_leave(self, original_node)
731 # Now, call any visitors that were hooked using a leave decorator.
732 for matcher, leave_funcs in reversed(list(self._extra_leave_funcs.items())):
733 if not self.matches(original_node, matcher):
734 continue
735 for leave_func in leave_funcs:
736 if _should_allow_visit(self._matchers, leave_func):
737 leave_func(original_node)
739 # Now, see if we have any matchers we should deactivate.
740 self._matchers = _leave_matchers(self._matchers, original_node)
742 def on_visit_attribute(self, node: cst.CSTNode, attribute: str) -> None:
743 # Evaluate whether this current function has a decorator on it.
744 if _should_allow_visit(
745 self._matchers,
746 getattr(self, f"visit_{type(node).__name__}_{attribute}", None),
747 ):
748 # Either the visit_func doesn't exist, we have no matchers, or we passed all
749 # matchers. In either case, just call the superclass behavior.
750 return CSTVisitor.on_visit_attribute(self, node, attribute)
752 def on_leave_attribute(self, original_node: cst.CSTNode, attribute: str) -> None:
753 # Evaluate whether this current function has a decorator on it.
754 if _should_allow_visit(
755 self._matchers,
756 getattr(self, f"leave_{type(original_node).__name__}_{attribute}", None),
757 ):
758 # Either the visit_func doesn't exist, we have no matchers, or we passed all
759 # matchers. In either case, just call the superclass behavior.
760 CSTVisitor.on_leave_attribute(self, original_node, attribute)
762 def matches(
763 self,
764 node: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
765 matcher: BaseMatcherNode,
766 ) -> bool:
767 """
768 A convenience method to call :func:`~libcst.matchers.matches` without requiring
769 an explicit parameter for metadata. Since our instance is an instance of
770 :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see
771 documentation for :func:`~libcst.matchers.matches` as it is identical to this
772 function.
773 """
774 return matches(node, matcher, metadata_resolver=self)
776 def findall(
777 self,
778 tree: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
779 matcher: Union[
780 BaseMatcherNode,
781 MatchIfTrue[cst.CSTNode],
782 MatchMetadata,
783 MatchMetadataIfTrue,
784 ],
785 ) -> Sequence[cst.CSTNode]:
786 """
787 A convenience method to call :func:`~libcst.matchers.findall` without requiring
788 an explicit parameter for metadata. Since our instance is an instance of
789 :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see
790 documentation for :func:`~libcst.matchers.findall` as it is identical to this
791 function.
792 """
793 return findall(tree, matcher, metadata_resolver=self)
795 def extract(
796 self,
797 node: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
798 matcher: BaseMatcherNode,
799 ) -> Optional[Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]]]:
800 """
801 A convenience method to call :func:`~libcst.matchers.extract` without requiring
802 an explicit parameter for metadata. Since our instance is an instance of
803 :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see
804 documentation for :func:`~libcst.matchers.extract` as it is identical to this
805 function.
806 """
807 return extract(node, matcher, metadata_resolver=self)
809 def extractall(
810 self,
811 tree: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
812 matcher: Union[
813 BaseMatcherNode,
814 MatchIfTrue[cst.CSTNode],
815 MatchMetadata,
816 MatchMetadataIfTrue,
817 ],
818 ) -> Sequence[Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]]]:
819 """
820 A convenience method to call :func:`~libcst.matchers.extractall` without requiring
821 an explicit parameter for metadata. Since our instance is an instance of
822 :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see
823 documentation for :func:`~libcst.matchers.extractall` as it is identical to this
824 function.
825 """
826 return extractall(tree, matcher, metadata_resolver=self)
828 def replace(
829 self,
830 tree: Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
831 matcher: Union[
832 BaseMatcherNode,
833 MatchIfTrue[cst.CSTNode],
834 MatchMetadata,
835 MatchMetadataIfTrue,
836 ],
837 replacement: Union[
838 cst.MaybeSentinel,
839 cst.RemovalSentinel,
840 cst.CSTNode,
841 Callable[
842 [cst.CSTNode, Dict[str, Union[cst.CSTNode, Sequence[cst.CSTNode]]]],
843 Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode],
844 ],
845 ],
846 ) -> Union[cst.MaybeSentinel, cst.RemovalSentinel, cst.CSTNode]:
847 """
848 A convenience method to call :func:`~libcst.matchers.replace` without requiring
849 an explicit parameter for metadata. Since our instance is an instance of
850 :class:`libcst.MetadataDependent`, we work as a metadata resolver. Please see
851 documentation for :func:`~libcst.matchers.replace` as it is identical to this
852 function.
853 """
854 return replace(tree, matcher, replacement, metadata_resolver=self)