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.
5
6import collections.abc
7import inspect
8import re
9from abc import ABCMeta
10from dataclasses import dataclass, fields
11from enum import auto, Enum
12from typing import (
13 Callable,
14 cast,
15 Dict,
16 Generic,
17 Iterator,
18 List,
19 Mapping,
20 NoReturn,
21 Optional,
22 Pattern,
23 Sequence,
24 Tuple,
25 Type,
26 TypeVar,
27 Union,
28)
29
30import libcst
31import libcst.metadata as meta
32from libcst import CSTLogicError, FlattenSentinel, MaybeSentinel, RemovalSentinel
33from libcst._metadata_dependent import LazyValue
34
35
36class DoNotCareSentinel(Enum):
37 """
38 A sentinel that is used in matcher classes to indicate that a caller
39 does not care what this value is. We recommend that you do not use this
40 directly, and instead use the :func:`DoNotCare` helper. You do not
41 need to use this for concrete matcher attributes since :func:`DoNotCare`
42 is already the default.
43 """
44
45 DEFAULT = auto()
46
47 def __repr__(self) -> str:
48 return "DoNotCare()"
49
50
51_MatcherT = TypeVar("_MatcherT", covariant=True)
52_MatchIfTrueT = TypeVar("_MatchIfTrueT", covariant=True)
53_BaseMatcherNodeSelfT = TypeVar("_BaseMatcherNodeSelfT", bound="BaseMatcherNode")
54_OtherNodeT = TypeVar("_OtherNodeT")
55_MetadataValueT = TypeVar("_MetadataValueT")
56_MatcherTypeT = TypeVar("_MatcherTypeT", bound=Type["BaseMatcherNode"])
57_OtherNodeMatcherTypeT = TypeVar(
58 "_OtherNodeMatcherTypeT", bound=Type["BaseMatcherNode"]
59)
60
61
62_METADATA_MISSING_SENTINEL = object()
63
64
65class AbstractBaseMatcherNodeMeta(ABCMeta):
66 """
67 Metaclass that all matcher nodes uses. Allows chaining 2 node type
68 together with an bitwise-or operator to produce an :class:`TypeOf`
69 matcher.
70 """
71
72 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
73 def __or__(self, node: Type["BaseMatcherNode"]) -> "TypeOf[Type[BaseMatcherNode]]":
74 return TypeOf(self, node)
75
76
77class BaseMatcherNode:
78 """
79 Base class that all concrete matchers subclass from. :class:`OneOf` and
80 :class:`AllOf` also subclass from this in order to allow them to be used in
81 any place that a concrete matcher is allowed. This means that, for example,
82 you can call :func:`matches` with a concrete matcher, or a :class:`OneOf` with
83 several concrete matchers as options.
84 """
85
86 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
87 def __or__(
88 self: _BaseMatcherNodeSelfT, other: _OtherNodeT
89 ) -> "OneOf[Union[_BaseMatcherNodeSelfT, _OtherNodeT]]":
90 return OneOf(self, other)
91
92 def __and__(
93 self: _BaseMatcherNodeSelfT, other: _OtherNodeT
94 ) -> "AllOf[Union[_BaseMatcherNodeSelfT, _OtherNodeT]]":
95 return AllOf(self, other)
96
97 def __invert__(self: _BaseMatcherNodeSelfT) -> "_BaseMatcherNodeSelfT":
98 return cast(_BaseMatcherNodeSelfT, _InverseOf(self))
99
100
101def DoNotCare() -> DoNotCareSentinel:
102 """
103 Used when you want to match exactly one node, but you do not care what node it is.
104 Useful inside sequences such as a :class:`libcst.matchers.Call`'s args attribte.
105 You do not need to use this for concrete matcher attributes since :func:`DoNotCare`
106 is already the default.
107
108 For example, the following matcher would match against any function calls with
109 three arguments, regardless of the arguments themselves and regardless of the
110 function name that we were calling::
111
112 m.Call(args=[m.DoNotCare(), m.DoNotCare(), m.DoNotCare()])
113 """
114 return DoNotCareSentinel.DEFAULT
115
116
117class TypeOf(Generic[_MatcherTypeT], BaseMatcherNode):
118 """
119 Matcher that matches any one of the given types. Useful when you want to work
120 with trees where a common property might belong to more than a single type.
121
122 For example, if you want either a binary operation or a boolean operation
123 where the left side has a name ``foo``::
124
125 m.TypeOf(m.BinaryOperation, m.BooleanOperation)(left = m.Name("foo"))
126
127 Or you could use the shorthand, like::
128
129 (m.BinaryOperation | m.BooleanOperation)(left = m.Name("foo"))
130
131 Also :class:`TypeOf` matchers can be used with initalizing in the default
132 state of other node matchers (without passing any extra patterns)::
133
134 m.Name | m.SimpleString
135
136 The will be equal to::
137
138 m.OneOf(m.Name(), m.SimpleString())
139 """
140
141 def __init__(self, *options: Union[_MatcherTypeT, "TypeOf[_MatcherTypeT]"]) -> None:
142 actual_options: List[_MatcherTypeT] = []
143 for option in options:
144 if isinstance(option, TypeOf):
145 if option.initalized:
146 raise ValueError(
147 "Cannot chain an uninitalized TypeOf with an initalized one"
148 )
149 actual_options.extend(option._raw_options)
150 else:
151 actual_options.append(option)
152
153 self._initalized = False
154 self._call_items: Tuple[Tuple[object, ...], Dict[str, object]] = ((), {})
155 self._raw_options: Tuple[_MatcherTypeT, ...] = tuple(actual_options)
156
157 @property
158 def initalized(self) -> bool:
159 return self._initalized
160
161 @property
162 def options(self) -> Iterator[BaseMatcherNode]:
163 for option in self._raw_options:
164 args, kwargs = self._call_items
165 matcher_pattern = option(*args, **kwargs)
166 yield matcher_pattern
167
168 def __call__(self, *args: object, **kwargs: object) -> BaseMatcherNode:
169 self._initalized = True
170 self._call_items = (args, kwargs)
171 return self
172
173 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
174 def __or__(
175 self, other: _OtherNodeMatcherTypeT
176 ) -> "TypeOf[Union[_MatcherTypeT, _OtherNodeMatcherTypeT]]":
177 return TypeOf[Union[_MatcherTypeT, _OtherNodeMatcherTypeT]](self, other)
178
179 # pyre-fixme[14]: `__and__` overrides method defined in `BaseMatcherNode`
180 # inconsistently.
181 def __and__(self, other: _OtherNodeMatcherTypeT) -> NoReturn:
182 left, right = type(self).__name__, other.__name__
183 raise TypeError(
184 f"TypeError: unsupported operand type(s) for &: {left!r} and {right!r}"
185 )
186
187 def __invert__(self) -> "AllOf[BaseMatcherNode]":
188 return AllOf(*map(DoesNotMatch, self.options))
189
190 def __repr__(self) -> str:
191 types = ", ".join(repr(option) for option in self._raw_options)
192 return f"TypeOf({types}, initalized = {self.initalized})"
193
194
195class OneOf(Generic[_MatcherT], BaseMatcherNode):
196 """
197 Matcher that matches any one of its options. Useful when you want to match
198 against one of several options for a single node. You can also construct a
199 :class:`OneOf` matcher by using Python's bitwise or operator with concrete
200 matcher classes.
201
202 For example, you could match against ``True``/``False`` like::
203
204 m.OneOf(m.Name("True"), m.Name("False"))
205
206 Or you could use the shorthand, like::
207
208 m.Name("True") | m.Name("False")
209
210 """
211
212 def __init__(self, *options: Union[_MatcherT, "OneOf[_MatcherT]"]) -> None:
213 actual_options: List[_MatcherT] = []
214 for option in options:
215 if isinstance(option, AllOf):
216 raise ValueError("Cannot use AllOf and OneOf in combination!")
217 elif isinstance(option, (OneOf, TypeOf)):
218 actual_options.extend(option.options)
219 else:
220 actual_options.append(option)
221 self._options: Sequence[_MatcherT] = tuple(actual_options)
222
223 @property
224 def options(self) -> Sequence[_MatcherT]:
225 """
226 The normalized list of options that we can choose from to satisfy a
227 :class:`OneOf` matcher. If any of these matchers are true, the
228 :class:`OneOf` matcher will also be considered a match.
229 """
230 return self._options
231
232 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
233 def __or__(self, other: _OtherNodeT) -> "OneOf[Union[_MatcherT, _OtherNodeT]]":
234 return OneOf(self, other)
235
236 def __and__(self, other: _OtherNodeT) -> NoReturn:
237 raise ValueError("Cannot use AllOf and OneOf in combination!")
238
239 def __invert__(self) -> "AllOf[_MatcherT]":
240 # Invert using De Morgan's Law so we don't have to complicate types.
241 return AllOf(*[DoesNotMatch(m) for m in self._options])
242
243 def __repr__(self) -> str:
244 return f"OneOf({', '.join([repr(o) for o in self._options])})"
245
246
247class AllOf(Generic[_MatcherT], BaseMatcherNode):
248 """
249 Matcher that matches all of its options. Useful when you want to match
250 against a concrete matcher and a :class:`MatchIfTrue` at the same time. Also
251 useful when you want to match against a concrete matcher and a
252 :func:`DoesNotMatch` at the same time. You can also construct a
253 :class:`AllOf` matcher by using Python's bitwise and operator with concrete
254 matcher classes.
255
256 For example, you could match against ``True`` in a roundabout way like::
257
258 m.AllOf(m.Name(), m.Name("True"))
259
260 Or you could use the shorthand, like::
261
262 m.Name() & m.Name("True")
263
264 Similar to :class:`OneOf`, this can be used in place of any concrete matcher.
265
266 Real-world cases where :class:`AllOf` is useful are hard to come by but they
267 are still provided for the limited edge cases in which they make sense. In
268 the example above, we are redundantly matching against any LibCST
269 :class:`~libcst.Name` node as well as LibCST :class:`~libcst.Name` nodes that
270 have the ``value`` of ``True``. We could drop the first option entirely and
271 get the same result. Often, if you are using a :class:`AllOf`,
272 you can refactor your code to be simpler.
273
274 For example, the following matches any function call to ``foo``, and
275 any function call which takes zero arguments::
276
277 m.AllOf(m.Call(func=m.Name("foo")), m.Call(args=()))
278
279 This could be refactored into the following equivalent concrete matcher::
280
281 m.Call(func=m.Name("foo"), args=())
282
283 """
284
285 def __init__(self, *options: Union[_MatcherT, "AllOf[_MatcherT]"]) -> None:
286 actual_options: List[_MatcherT] = []
287 for option in options:
288 if isinstance(option, OneOf):
289 raise ValueError("Cannot use AllOf and OneOf in combination!")
290 elif isinstance(option, TypeOf):
291 raise ValueError("Cannot use AllOf and TypeOf in combination!")
292 elif isinstance(option, AllOf):
293 actual_options.extend(option.options)
294 else:
295 actual_options.append(option)
296 self._options: Sequence[_MatcherT] = tuple(actual_options)
297
298 @property
299 def options(self) -> Sequence[_MatcherT]:
300 """
301 The normalized list of options that we can choose from to satisfy a
302 :class:`AllOf` matcher. If all of these matchers are true, the
303 :class:`AllOf` matcher will also be considered a match.
304 """
305 return self._options
306
307 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
308 def __or__(self, other: _OtherNodeT) -> NoReturn:
309 raise ValueError("Cannot use AllOf and OneOf in combination!")
310
311 def __and__(self, other: _OtherNodeT) -> "AllOf[Union[_MatcherT, _OtherNodeT]]":
312 return AllOf(self, other)
313
314 def __invert__(self) -> "OneOf[_MatcherT]":
315 # Invert using De Morgan's Law so we don't have to complicate types.
316 return OneOf(*[DoesNotMatch(m) for m in self._options])
317
318 def __repr__(self) -> str:
319 return f"AllOf({', '.join([repr(o) for o in self._options])})"
320
321
322class _InverseOf(Generic[_MatcherT]):
323 """
324 Matcher that inverts the match result of its child. You can also construct a
325 :class:`_InverseOf` matcher by using Python's bitwise invert operator with concrete
326 matcher classes or any special matcher.
327
328 Note that you should refrain from constructing a :class:`_InverseOf` directly, and
329 should instead use the :func:`DoesNotMatch` helper function.
330
331 For example, the following matches against any identifier that isn't
332 ``True``/``False``::
333
334 m.DoesNotMatch(m.OneOf(m.Name("True"), m.Name("False")))
335
336 Or you could use the shorthand, like:
337
338 ~(m.Name("True") | m.Name("False"))
339
340 """
341
342 def __init__(self, matcher: _MatcherT) -> None:
343 self._matcher: _MatcherT = matcher
344
345 @property
346 def matcher(self) -> _MatcherT:
347 """
348 The matcher that we will evaluate and invert. If this matcher is true, then
349 :class:`_InverseOf` will be considered not a match, and vice-versa.
350 """
351 return self._matcher
352
353 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
354 def __or__(self, other: _OtherNodeT) -> "OneOf[Union[_MatcherT, _OtherNodeT]]":
355 # Without a cast, pyre thinks that the below OneOf is type OneOf[object]
356 # even though it has the types passed into it.
357 return cast(OneOf[Union[_MatcherT, _OtherNodeT]], OneOf(self, other))
358
359 def __and__(self, other: _OtherNodeT) -> "AllOf[Union[_MatcherT, _OtherNodeT]]":
360 # Without a cast, pyre thinks that the below AllOf is type AllOf[object]
361 # even though it has the types passed into it.
362 return cast(AllOf[Union[_MatcherT, _OtherNodeT]], AllOf(self, other))
363
364 def __getattr__(self, key: str) -> object:
365 # We lie about types to make _InverseOf appear transparent. So, its conceivable
366 # that somebody might try to dereference an attribute on the _MatcherT wrapped
367 # node and become surprised that it doesn't work.
368 return getattr(self._matcher, key)
369
370 def __invert__(self) -> _MatcherT:
371 return self._matcher
372
373 def __repr__(self) -> str:
374 return f"DoesNotMatch({repr(self._matcher)})"
375
376
377class _ExtractMatchingNode(Generic[_MatcherT]):
378 """
379 Transparent pass-through matcher that captures the node which matches its children,
380 making it available to the caller of :func:`extract` or :func:`extractall`.
381
382 Note that you should refrain from constructing a :class:`_ExtractMatchingNode`
383 directly, and should instead use the :func:`SaveMatchedNode` helper function.
384
385 For example, the following will match against any binary operation whose left
386 and right operands are not integers, saving those expressions for later inspection.
387 If used inside :func:`extract` or :func:`extractall`, the resulting dictionary will
388 contain the keys ``left_operand`` and ``right_operand``.
389
390 m.BinaryOperation(
391 left=m.SaveMatchedNode(
392 m.DoesNotMatch(m.Integer()),
393 "left_operand",
394 ),
395 right=m.SaveMatchedNode(
396 m.DoesNotMatch(m.Integer()),
397 "right_operand",
398 ),
399 )
400 """
401
402 def __init__(self, matcher: _MatcherT, name: str) -> None:
403 self._matcher: _MatcherT = matcher
404 self._name: str = name
405
406 @property
407 def matcher(self) -> _MatcherT:
408 """
409 The matcher that we will evaluate and capture matching LibCST nodes for.
410 If this matcher is true, then :class:`_ExtractMatchingNode` will be considered
411 a match and will save the node which matched.
412 """
413 return self._matcher
414
415 @property
416 def name(self) -> str:
417 """
418 The name we will call our captured LibCST node inside the resulting dictionary
419 returned by :func:`extract` or :func:`extractall`.
420 """
421 return self._name
422
423 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
424 def __or__(self, other: _OtherNodeT) -> "OneOf[Union[_MatcherT, _OtherNodeT]]":
425 # Without a cast, pyre thinks that the below OneOf is type OneOf[object]
426 # even though it has the types passed into it.
427 return cast(OneOf[Union[_MatcherT, _OtherNodeT]], OneOf(self, other))
428
429 def __and__(self, other: _OtherNodeT) -> "AllOf[Union[_MatcherT, _OtherNodeT]]":
430 # This doesn't make sense. If we have multiple SaveMatchedNode captures
431 # that are captured with an and, either all of them will be assigned the
432 # same node, or none of them. It makes more sense to move the SaveMatchedNode
433 # up to wrap the AllOf.
434 raise ValueError(
435 (
436 "Cannot use AllOf with SavedMatchedNode children! Instead, you should "
437 + "use SaveMatchedNode(AllOf(options...))."
438 )
439 )
440
441 def __getattr__(self, key: str) -> object:
442 # We lie about types to make _ExtractMatchingNode appear transparent. So,
443 # its conceivable that somebody might try to dereference an attribute on
444 # the _MatcherT wrapped node and become surprised that it doesn't work.
445 return getattr(self._matcher, key)
446
447 def __invert__(self) -> "_MatcherT":
448 # This doesn't make sense. We don't want to capture a node only if it
449 # doesn't match, since this will never capture anything.
450 raise ValueError(
451 (
452 "Cannot invert a SaveMatchedNode. Instead you should wrap SaveMatchedNode "
453 "around your inversion itself"
454 )
455 )
456
457 def __repr__(self) -> str:
458 return (
459 f"SaveMatchedNode(matcher={repr(self._matcher)}, name={repr(self._name)})"
460 )
461
462
463class MatchIfTrue(Generic[_MatchIfTrueT]):
464 """
465 Matcher that matches if its child callable returns ``True``. The child callable
466 should take one argument which is the attribute on the LibCST node we are
467 trying to match against. This is useful if you want to do complex logic to
468 determine if an attribute should match or not. One example of this is the
469 :func:`MatchRegex` matcher build on top of :class:`MatchIfTrue` which takes a
470 regular expression and matches any string attribute where a regex match is found.
471
472 For example, to match on any identifier spelled with the letter ``e``::
473
474 m.Name(value=m.MatchIfTrue(lambda value: "e" in value))
475
476 This can be used in place of any concrete matcher as long as it is not the
477 root matcher. Calling :func:`matches` directly on a :class:`MatchIfTrue` is
478 redundant since you can just call the child callable directly with the node
479 you are passing to :func:`matches`.
480 """
481
482 _func: Callable[[_MatchIfTrueT], bool]
483
484 def __init__(self, func: Callable[[_MatchIfTrueT], bool]) -> None:
485 self._func = func
486
487 @property
488 def func(self) -> Callable[[_MatchIfTrueT], bool]:
489 """
490 The function that we will call with a LibCST node in order to determine
491 if we match. If the function returns ``True`` then we consider ourselves
492 to be a match.
493 """
494 return self._func
495
496 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
497 def __or__(
498 self, other: _OtherNodeT
499 ) -> "OneOf[Union[MatchIfTrue[_MatchIfTrueT], _OtherNodeT]]":
500 return OneOf(self, other)
501
502 def __and__(
503 self, other: _OtherNodeT
504 ) -> "AllOf[Union[MatchIfTrue[_MatchIfTrueT], _OtherNodeT]]":
505 return AllOf(self, other)
506
507 def __invert__(self) -> "MatchIfTrue[_MatchIfTrueT]":
508 # Construct a wrapped version of MatchIfTrue for typing simplicity.
509 # Without the cast, pyre doesn't seem to think the lambda is valid.
510 return MatchIfTrue(lambda val: not self._func(val))
511
512 def __repr__(self) -> str:
513 return f"MatchIfTrue({repr(self._func)})"
514
515
516def MatchRegex(regex: Union[str, Pattern[str]]) -> MatchIfTrue[str]:
517 """
518 Used as a convenience wrapper to :class:`MatchIfTrue` which allows for
519 matching a string attribute against a regex. ``regex`` can be any regular
520 expression string or a compiled ``Pattern``. This uses Python's re module
521 under the hood and is compatible with syntax documented on
522 `docs.python.org <https://docs.python.org/3/library/re.html>`_.
523
524 For example, to match against any identifier that is at least one character
525 long and only contains alphabetical characters::
526
527 m.Name(value=m.MatchRegex(r'[A-Za-z]+'))
528
529 This can be used in place of any string literal when constructing a concrete
530 matcher.
531 """
532
533 def _match_func(value: object) -> bool:
534 if isinstance(value, str):
535 return bool(re.fullmatch(regex, value))
536 else:
537 return False
538
539 return MatchIfTrue(_match_func)
540
541
542class _BaseMetadataMatcher:
543 """
544 Class that's only around for typing purposes.
545 """
546
547 pass
548
549
550class MatchMetadata(_BaseMetadataMatcher):
551 """
552 Matcher that looks up the metadata on the current node using the provided
553 metadata provider and compares the value on the node against the value provided
554 to :class:`MatchMetadata`.
555 If the metadata provider is unresolved, a :class:`LookupError` exeption will be
556 raised and ask you to provide a :class:`~libcst.metadata.MetadataWrapper`.
557 If the metadata value does not exist for a particular node, :class:`MatchMetadata`
558 will be considered not a match.
559
560 For example, to match against any function call which has one parameter which
561 is used in a load expression context::
562
563 m.Call(
564 args=[
565 m.Arg(
566 m.MatchMetadata(
567 meta.ExpressionContextProvider,
568 meta.ExpressionContext.LOAD,
569 )
570 )
571 ]
572 )
573
574 To match against any :class:`~libcst.Name` node for the identifier ``foo``
575 which is the target of an assignment::
576
577 m.Name(
578 value="foo",
579 metadata=m.MatchMetadata(
580 meta.ExpressionContextProvider,
581 meta.ExpressionContext.STORE,
582 )
583 )
584
585 This can be used in place of any concrete matcher as long as it is not the
586 root matcher. Calling :func:`matches` directly on a :class:`MatchMetadata` is
587 redundant since you can just check the metadata on the root node that you
588 are passing to :func:`matches`.
589 """
590
591 def __init__(
592 self,
593 key: Type[meta.BaseMetadataProvider[_MetadataValueT]],
594 value: _MetadataValueT,
595 ) -> None:
596 self._key: Type[meta.BaseMetadataProvider[_MetadataValueT]] = key
597 self._value: _MetadataValueT = value
598
599 @property
600 def key(self) -> meta.ProviderT:
601 """
602 The metadata provider that we will use to fetch values when identifying whether
603 a node matches this matcher. We compare the value returned from the metadata
604 provider to the value provided in ``value`` when determining a match.
605 """
606 return self._key
607
608 @property
609 def value(self) -> object:
610 """
611 The value that we will compare against the return from the metadata provider
612 for each node when determining a match.
613 """
614 return self._value
615
616 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
617 def __or__(self, other: _OtherNodeT) -> "OneOf[Union[MatchMetadata, _OtherNodeT]]":
618 return OneOf(self, other)
619
620 def __and__(self, other: _OtherNodeT) -> "AllOf[Union[MatchMetadata, _OtherNodeT]]":
621 return AllOf(self, other)
622
623 def __invert__(self) -> "MatchMetadata":
624 # We intentionally lie here, for the same reason given in the documentation
625 # for DoesNotMatch.
626 return cast(MatchMetadata, _InverseOf(self))
627
628 def __repr__(self) -> str:
629 return f"MatchMetadata(key={repr(self._key)}, value={repr(self._value)})"
630
631
632class MatchMetadataIfTrue(_BaseMetadataMatcher):
633 """
634 Matcher that looks up the metadata on the current node using the provided
635 metadata provider and passes it to a callable which can inspect the metadata
636 further, returning ``True`` if the matcher should be considered a match.
637 If the metadata provider is unresolved, a :class:`LookupError` exeption will be
638 raised and ask you to provide a :class:`~libcst.metadata.MetadataWrapper`.
639 If the metadata value does not exist for a particular node,
640 :class:`MatchMetadataIfTrue` will be considered not a match.
641
642 For example, to match against any arg whose qualified name might be
643 ``typing.Dict``::
644
645 m.Call(
646 args=[
647 m.Arg(
648 m.MatchMetadataIfTrue(
649 meta.QualifiedNameProvider,
650 lambda qualnames: any(n.name == "typing.Dict" for n in qualnames)
651 )
652 )
653 ]
654 )
655
656 To match against any :class:`~libcst.Name` node for the identifier ``foo``
657 as long as that identifier is found at the beginning of an unindented line::
658
659 m.Name(
660 value="foo",
661 metadata=m.MatchMetadataIfTrue(
662 meta.PositionProvider,
663 lambda position: position.start.column == 0,
664 )
665 )
666
667 This can be used in place of any concrete matcher as long as it is not the
668 root matcher. Calling :func:`matches` directly on a :class:`MatchMetadataIfTrue`
669 is redundant since you can just check the metadata on the root node that you
670 are passing to :func:`matches`.
671 """
672
673 def __init__(
674 self,
675 key: Type[meta.BaseMetadataProvider[_MetadataValueT]],
676 func: Callable[[_MetadataValueT], bool],
677 ) -> None:
678 self._key: Type[meta.BaseMetadataProvider[_MetadataValueT]] = key
679 self._func: Callable[[_MetadataValueT], bool] = func
680
681 @property
682 def key(self) -> meta.ProviderT:
683 """
684 The metadata provider that we will use to fetch values when identifying whether
685 a node matches this matcher. We pass the value returned from the metadata
686 provider to the callable given to us in ``func``.
687 """
688 return self._key
689
690 @property
691 def func(self) -> Callable[[object], bool]:
692 """
693 The function that we will call with a value retrieved from the metadata provider
694 provided in ``key``. If the function returns ``True`` then we consider ourselves
695 to be a match.
696 """
697 return self._func
698
699 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
700 def __or__(
701 self, other: _OtherNodeT
702 ) -> "OneOf[Union[MatchMetadataIfTrue, _OtherNodeT]]":
703 return OneOf(self, other)
704
705 def __and__(
706 self, other: _OtherNodeT
707 ) -> "AllOf[Union[MatchMetadataIfTrue, _OtherNodeT]]":
708 return AllOf(self, other)
709
710 def __invert__(self) -> "MatchMetadataIfTrue":
711 # Construct a wrapped version of MatchMetadataIfTrue for typing simplicity.
712 return MatchMetadataIfTrue(self._key, lambda val: not self._func(val))
713
714 def __repr__(self) -> str:
715 return f"MatchMetadataIfTrue(key={repr(self._key)}, func={repr(self._func)})"
716
717
718class _BaseWildcardNode:
719 """
720 A typing-only class for internal helpers in this module to be able to
721 specify that they take a wildcard node type.
722 """
723
724 pass
725
726
727class AtLeastN(Generic[_MatcherT], _BaseWildcardNode):
728 """
729 Matcher that matches ``n`` or more LibCST nodes in a row in a sequence.
730 :class:`AtLeastN` defaults to matching against the :func:`DoNotCare` matcher,
731 so if you do not specify a matcher as a child, :class:`AtLeastN`
732 will match only by count. If you do specify a matcher as a child,
733 :class:`AtLeastN` will instead make sure that each LibCST node matches the
734 matcher supplied.
735
736 For example, this will match all function calls with at least 3 arguments::
737
738 m.Call(args=[m.AtLeastN(n=3)])
739
740 This will match all function calls with 3 or more integer arguments::
741
742 m.Call(args=[m.AtLeastN(n=3, matcher=m.Arg(m.Integer()))])
743
744 You can combine sequence matchers with concrete matchers and special matchers
745 and it will behave as you expect. For example, this will match all function
746 calls that have 2 or more integer arguments in a row, followed by any arbitrary
747 argument::
748
749 m.Call(args=[m.AtLeastN(n=2, matcher=m.Arg(m.Integer())), m.DoNotCare()])
750
751 And finally, this will match all function calls that have at least 5
752 arguments, the final one being an integer::
753
754 m.Call(args=[m.AtLeastN(n=4), m.Arg(m.Integer())])
755 """
756
757 def __init__(
758 self,
759 matcher: Union[_MatcherT, DoNotCareSentinel] = DoNotCareSentinel.DEFAULT,
760 *,
761 n: int,
762 ) -> None:
763 if n < 0:
764 raise ValueError(
765 f"{self.__class__.__qualname__} n attribute must be positive"
766 )
767 self._n: int = n
768 self._matcher: Union[_MatcherT, DoNotCareSentinel] = matcher
769
770 @property
771 def n(self) -> int:
772 """
773 The number of nodes in a row that must match :attr:`AtLeastN.matcher` for
774 this matcher to be considered a match. If there are less than ``n`` matches,
775 this matcher will not be considered a match. If there are equal to or more
776 than ``n`` matches, this matcher will be considered a match.
777 """
778 return self._n
779
780 @property
781 def matcher(self) -> Union[_MatcherT, DoNotCareSentinel]:
782 """
783 The matcher which each node in a sequence needs to match.
784 """
785 return self._matcher
786
787 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
788 def __or__(self, other: object) -> NoReturn:
789 raise ValueError("AtLeastN cannot be used in a OneOf matcher")
790
791 def __and__(self, other: object) -> NoReturn:
792 raise ValueError("AtLeastN cannot be used in an AllOf matcher")
793
794 def __invert__(self) -> NoReturn:
795 raise ValueError("Cannot invert an AtLeastN matcher!")
796
797 def __repr__(self) -> str:
798 if self._n == 0:
799 return f"ZeroOrMore({repr(self._matcher)})"
800 else:
801 return f"AtLeastN({repr(self._matcher)}, n={self._n})"
802
803
804def ZeroOrMore(
805 matcher: Union[_MatcherT, DoNotCareSentinel] = DoNotCareSentinel.DEFAULT,
806) -> AtLeastN[Union[_MatcherT, DoNotCareSentinel]]:
807 """
808 Used as a convenience wrapper to :class:`AtLeastN` when ``n`` is equal to ``0``.
809 Use this when you want to match against any number of nodes in a sequence.
810
811 For example, this will match any function call with zero or more arguments, as
812 long as all of the arguments are integers::
813
814 m.Call(args=[m.ZeroOrMore(m.Arg(m.Integer()))])
815
816 This will match any function call where the first argument is an integer and
817 it doesn't matter what the rest of the arguments are::
818
819 m.Call(args=[m.Arg(m.Integer()), m.ZeroOrMore()])
820
821 You will often want to use :class:`ZeroOrMore` on both sides of a concrete
822 matcher in order to match against sequences that contain a particular node
823 in an arbitrary location. For example, the following will match any function
824 call that takes in at least one string argument anywhere::
825
826 m.Call(args=[m.ZeroOrMore(), m.Arg(m.SimpleString()), m.ZeroOrMore()])
827 """
828 return cast(AtLeastN[Union[_MatcherT, DoNotCareSentinel]], AtLeastN(matcher, n=0))
829
830
831class AtMostN(Generic[_MatcherT], _BaseWildcardNode):
832 """
833 Matcher that matches ``n`` or fewer LibCST nodes in a row in a sequence.
834 :class:`AtMostN` defaults to matching against the :func:`DoNotCare` matcher,
835 so if you do not specify a matcher as a child, :class:`AtMostN` will
836 match only by count. If you do specify a matcher as a child,
837 :class:`AtMostN` will instead make sure that each LibCST node matches the
838 matcher supplied.
839
840 For example, this will match all function calls with 3 or fewer arguments::
841
842 m.Call(args=[m.AtMostN(n=3)])
843
844 This will match all function calls with 0, 1 or 2 string arguments::
845
846 m.Call(args=[m.AtMostN(n=2, matcher=m.Arg(m.SimpleString()))])
847
848 You can combine sequence matchers with concrete matchers and special matchers
849 and it will behave as you expect. For example, this will match all function
850 calls that have 0, 1 or 2 string arguments in a row, followed by an arbitrary
851 argument::
852
853 m.Call(args=[m.AtMostN(n=2, matcher=m.Arg(m.SimpleString())), m.DoNotCare()])
854
855 And finally, this will match all function calls that have at least 2
856 arguments, the final one being a string::
857
858 m.Call(args=[m.AtMostN(n=2), m.Arg(m.SimpleString())])
859 """
860
861 def __init__(
862 self,
863 matcher: Union[_MatcherT, DoNotCareSentinel] = DoNotCareSentinel.DEFAULT,
864 *,
865 n: int,
866 ) -> None:
867 if n < 0:
868 raise ValueError(
869 f"{self.__class__.__qualname__} n attribute must be positive"
870 )
871 self._n: int = n
872 self._matcher: Union[_MatcherT, DoNotCareSentinel] = matcher
873
874 @property
875 def n(self) -> int:
876 """
877 The number of nodes in a row that must match :attr:`AtLeastN.matcher` for
878 this matcher to be considered a match. If there are less than or equal to
879 ``n`` matches, then this matcher will be considered a match. Any more than
880 ``n`` matches in a row and this matcher will stop matching and be considered
881 not a match.
882 """
883 return self._n
884
885 @property
886 def matcher(self) -> Union[_MatcherT, DoNotCareSentinel]:
887 """
888 The matcher which each node in a sequence needs to match.
889 """
890 return self._matcher
891
892 # pyre-fixme[15]: `__or__` overrides method defined in `type` inconsistently.
893 def __or__(self, other: object) -> NoReturn:
894 raise ValueError("AtMostN cannot be used in a OneOf matcher")
895
896 def __and__(self, other: object) -> NoReturn:
897 raise ValueError("AtMostN cannot be used in an AllOf matcher")
898
899 def __invert__(self) -> NoReturn:
900 raise ValueError("Cannot invert an AtMostN matcher!")
901
902 def __repr__(self) -> str:
903 if self._n == 1:
904 return f"ZeroOrOne({repr(self._matcher)})"
905 else:
906 return f"AtMostN({repr(self._matcher)}, n={self._n})"
907
908
909def ZeroOrOne(
910 matcher: Union[_MatcherT, DoNotCareSentinel] = DoNotCareSentinel.DEFAULT,
911) -> AtMostN[Union[_MatcherT, DoNotCareSentinel]]:
912 """
913 Used as a convenience wrapper to :class:`AtMostN` when ``n`` is equal to ``1``.
914 This is effectively a maybe clause.
915
916 For example, this will match any function call with zero or one integer
917 argument::
918
919 m.Call(args=[m.ZeroOrOne(m.Arg(m.Integer()))])
920
921 This will match any function call that has two or three arguments, and
922 the first and last arguments are strings::
923
924 m.Call(args=[m.Arg(m.SimpleString()), m.ZeroOrOne(), m.Arg(m.SimpleString())])
925
926 """
927 return cast(AtMostN[Union[_MatcherT, DoNotCareSentinel]], AtMostN(matcher, n=1))
928
929
930def DoesNotMatch(obj: _OtherNodeT) -> _OtherNodeT:
931 """
932 Matcher helper that inverts the match result of its child. You can also invert a
933 matcher by using Python's bitwise invert operator on concrete matchers or any
934 special matcher.
935
936 For example, the following matches against any identifier that isn't
937 ``True``/``False``::
938
939 m.DoesNotMatch(m.OneOf(m.Name("True"), m.Name("False")))
940
941 Or you could use the shorthand, like::
942
943 ~(m.Name("True") | m.Name("False"))
944
945 This can be used in place of any concrete matcher as long as it is not the
946 root matcher. Calling :func:`matches` directly on a :func:`DoesNotMatch` is
947 redundant since you can invert the return of :func:`matches` using a bitwise not.
948 """
949
950 # This type is a complete, dirty lie, but there's no way to recursively apply
951 # a parameter to each type inside a Union that may be in a _OtherNodeT.
952 # However, given the way _InverseOf works (it will unwrap itself if
953 # inverted again), and the way we apply De Morgan's law for OneOf and AllOf,
954 # this lie ends up getting us correct typing. Anywhere a node is valid, using
955 # DoesNotMatch(node) is also valid.
956 #
957 # ~MatchIfTrue is still MatchIfTrue
958 # ~MatchMetadataIfTrue is still MatchMetadataIfTrue
959 # ~OneOf[x] is AllOf[~x]
960 # ~AllOf[x] is OneOf[~x]
961 # ~~x is x
962 #
963 # So, under all circumstances, since OneOf/AllOf are both allowed in every
964 # instance, and given that inverting MatchIfTrue is still MatchIfTrue,
965 # and inverting an inverted value returns us the original, its clear that
966 # there are no operations we can possibly do that bring us outside of the
967 # types specified in the concrete matchers as long as we lie that DoesNotMatch
968 # returns the value passed in.
969 if isinstance(
970 obj,
971 (
972 BaseMatcherNode,
973 MatchIfTrue,
974 _BaseMetadataMatcher,
975 _InverseOf,
976 _ExtractMatchingNode,
977 ),
978 ):
979 # We can use the overridden __invert__ in this case. Pyre doesn't think
980 # we can though, and casting doesn't fix the issue.
981 inverse = ~obj
982 else:
983 # We must wrap in a _InverseOf.
984 inverse = _InverseOf(obj)
985 return cast(_OtherNodeT, inverse)
986
987
988def SaveMatchedNode(matcher: _OtherNodeT, name: str) -> _OtherNodeT:
989 """
990 Matcher helper that captures the matched node that matched against a matcher
991 class, making it available in the dictionary returned by :func:`extract` or
992 :func:`extractall`.
993
994 For example, the following will match against any binary operation whose left
995 and right operands are not integers, saving those expressions for later inspection.
996 If used inside :func:`extract` or :func:`extractall`, the resulting dictionary
997 will contain the keys ``left_operand`` and ``right_operand``::
998
999 m.BinaryOperation(
1000 left=m.SaveMatchedNode(
1001 m.DoesNotMatch(m.Integer()),
1002 "left_operand",
1003 ),
1004 right=m.SaveMatchedNode(
1005 m.DoesNotMatch(m.Integer()),
1006 "right_operand",
1007 ),
1008 )
1009
1010 This can be used in place of any concrete matcher as long as it is not the
1011 root matcher. Calling :func:`extract` directly on a :func:`SaveMatchedNode` is
1012 redundant since you already have the reference to the node itself.
1013 """
1014 return cast(_OtherNodeT, _ExtractMatchingNode(matcher, name))
1015
1016
1017def _matches_zero_nodes(
1018 matcher: Union[
1019 BaseMatcherNode,
1020 _BaseWildcardNode,
1021 MatchIfTrue[libcst.CSTNode],
1022 _BaseMetadataMatcher,
1023 DoNotCareSentinel,
1024 ],
1025) -> bool:
1026 if isinstance(matcher, AtLeastN) and matcher.n == 0:
1027 return True
1028 if isinstance(matcher, AtMostN):
1029 return True
1030 if isinstance(matcher, _ExtractMatchingNode):
1031 return _matches_zero_nodes(matcher.matcher)
1032 return False
1033
1034
1035@dataclass(frozen=True)
1036class _SequenceMatchesResult:
1037 sequence_capture: Optional[
1038 Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]
1039 ]
1040 matched_nodes: Optional[
1041 Union[libcst.CSTNode, MaybeSentinel, Sequence[libcst.CSTNode]]
1042 ]
1043
1044
1045def _sequence_matches( # noqa: C901
1046 nodes: Sequence[Union[MaybeSentinel, libcst.CSTNode]],
1047 matchers: Sequence[
1048 Union[
1049 BaseMatcherNode,
1050 _BaseWildcardNode,
1051 MatchIfTrue[libcst.CSTNode],
1052 _BaseMetadataMatcher,
1053 DoNotCareSentinel,
1054 ]
1055 ],
1056 metadata_lookup: Callable[[meta.ProviderT, libcst.CSTNode], object],
1057) -> _SequenceMatchesResult:
1058 if not nodes and not matchers:
1059 # Base case, empty lists are always matches
1060 return _SequenceMatchesResult({}, None)
1061 if not nodes and matchers:
1062 # Base case, we have one or more matcher that wasn't matched
1063 if all(_matches_zero_nodes(m) for m in matchers):
1064 return _SequenceMatchesResult(
1065 # pyre-ignore[16]: `MatchIfTrue` has no attribute `name`.
1066 {m.name: () for m in matchers if isinstance(m, _ExtractMatchingNode)},
1067 (),
1068 )
1069 else:
1070 return _SequenceMatchesResult(None, None)
1071 if nodes and not matchers:
1072 # Base case, we have nodes left that don't match any matcher
1073 return _SequenceMatchesResult(None, None)
1074
1075 # Recursive case, nodes and matchers LHS matches
1076 node = nodes[0]
1077 matcher = matchers[0]
1078 if isinstance(matcher, DoNotCareSentinel):
1079 # We don't care about the value for this node.
1080 return _SequenceMatchesResult(
1081 _sequence_matches(
1082 nodes[1:], matchers[1:], metadata_lookup
1083 ).sequence_capture,
1084 node,
1085 )
1086 elif isinstance(matcher, _BaseWildcardNode):
1087 if isinstance(matcher, AtMostN):
1088 if matcher.n > 0:
1089 # First, assume that this does match a node (greedy).
1090 # Consume one node since it matched this matcher.
1091 attribute_capture = _attribute_matches(
1092 nodes[0], matcher.matcher, metadata_lookup
1093 )
1094 if attribute_capture is not None:
1095 result = _sequence_matches(
1096 nodes[1:],
1097 [AtMostN(matcher.matcher, n=matcher.n - 1), *matchers[1:]],
1098 metadata_lookup,
1099 )
1100 if result.sequence_capture is not None:
1101 matched = result.matched_nodes
1102 assert isinstance(matched, Sequence)
1103 return _SequenceMatchesResult(
1104 {**attribute_capture, **result.sequence_capture},
1105 # pyre-fixme[6]: Expected `Union[None, Sequence[libcst._n...
1106 (node, *matched),
1107 )
1108 # Finally, assume that this does not match the current node.
1109 # Consume the matcher but not the node.
1110 return _SequenceMatchesResult(
1111 _sequence_matches(
1112 nodes, matchers[1:], metadata_lookup
1113 ).sequence_capture,
1114 (),
1115 )
1116 elif isinstance(matcher, AtLeastN):
1117 if matcher.n > 0:
1118 # Only match if we can consume one of the matches, since we still
1119 # need to match N nodes.
1120 attribute_capture = _attribute_matches(
1121 nodes[0], matcher.matcher, metadata_lookup
1122 )
1123 if attribute_capture is not None:
1124 result = _sequence_matches(
1125 nodes[1:],
1126 [AtLeastN(matcher.matcher, n=matcher.n - 1), *matchers[1:]],
1127 metadata_lookup,
1128 )
1129 if result.sequence_capture is not None:
1130 matched = result.matched_nodes
1131 assert isinstance(matched, Sequence)
1132 return _SequenceMatchesResult(
1133 {**attribute_capture, **result.sequence_capture},
1134 # pyre-fixme[6]: Expected `Union[None, Sequence[libcst._n...
1135 (node, *matched),
1136 )
1137 return _SequenceMatchesResult(None, None)
1138 else:
1139 # First, assume that this does match a node (greedy).
1140 # Consume one node since it matched this matcher.
1141 attribute_capture = _attribute_matches(
1142 nodes[0], matcher.matcher, metadata_lookup
1143 )
1144 if attribute_capture is not None:
1145 result = _sequence_matches(nodes[1:], matchers, metadata_lookup)
1146 if result.sequence_capture is not None:
1147 matched = result.matched_nodes
1148 assert isinstance(matched, Sequence)
1149 return _SequenceMatchesResult(
1150 {**attribute_capture, **result.sequence_capture},
1151 # pyre-fixme[6]: Expected `Union[None, Sequence[libcst._n...
1152 (node, *matched),
1153 )
1154 # Now, assume that this does not match the current node.
1155 # Consume the matcher but not the node.
1156 return _SequenceMatchesResult(
1157 _sequence_matches(
1158 nodes, matchers[1:], metadata_lookup
1159 ).sequence_capture,
1160 (),
1161 )
1162 else:
1163 # There are no other types of wildcard consumers, but we're making
1164 # pyre happy with that fact.
1165 raise CSTLogicError(f"Logic error unrecognized wildcard {type(matcher)}!")
1166 elif isinstance(matcher, _ExtractMatchingNode):
1167 # See if the raw matcher matches. If it does, capture the sequence we matched and store it.
1168 result = _sequence_matches(
1169 nodes, [matcher.matcher, *matchers[1:]], metadata_lookup
1170 )
1171 if result.sequence_capture is not None:
1172 return _SequenceMatchesResult(
1173 {
1174 # Our own match capture comes first, since we wnat to allow the same
1175 # name later in the sequence to override us.
1176 matcher.name: result.matched_nodes,
1177 **result.sequence_capture,
1178 },
1179 result.matched_nodes,
1180 )
1181 return _SequenceMatchesResult(None, None)
1182
1183 match_capture = _matches(node, matcher, metadata_lookup)
1184 if match_capture is not None:
1185 # These values match directly
1186 result = _sequence_matches(nodes[1:], matchers[1:], metadata_lookup)
1187 if result.sequence_capture is not None:
1188 return _SequenceMatchesResult(
1189 {**match_capture, **result.sequence_capture}, node
1190 )
1191
1192 # Failed recursive case, no match
1193 return _SequenceMatchesResult(None, None)
1194
1195
1196_AttributeValueT = Optional[Union[MaybeSentinel, libcst.CSTNode, str, bool]]
1197_AttributeMatcherT = Optional[Union[BaseMatcherNode, DoNotCareSentinel, str, bool]]
1198
1199
1200def _attribute_matches( # noqa: C901
1201 node: Union[_AttributeValueT, Sequence[_AttributeValueT]],
1202 matcher: Union[_AttributeMatcherT, Sequence[_AttributeMatcherT]],
1203 metadata_lookup: Callable[[meta.ProviderT, libcst.CSTNode], object],
1204) -> Optional[Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]]:
1205 if isinstance(matcher, DoNotCareSentinel):
1206 # We don't care what this is, so don't penalize a non-match.
1207 return {}
1208 if isinstance(matcher, _InverseOf):
1209 # Return the opposite evaluation
1210 return (
1211 {}
1212 if _attribute_matches(node, matcher.matcher, metadata_lookup) is None
1213 else None
1214 )
1215 if isinstance(matcher, _ExtractMatchingNode):
1216 attribute_capture = _attribute_matches(node, matcher.matcher, metadata_lookup)
1217 if attribute_capture is not None:
1218 return {
1219 # Our own match capture comes last, since its higher in the tree
1220 # so we want to override any child match captures by the same name.
1221 **attribute_capture,
1222 matcher.name: node,
1223 }
1224 return None
1225
1226 if isinstance(matcher, MatchIfTrue):
1227 # We should only return if the matcher function is true.
1228 return {} if matcher.func(node) else None
1229
1230 if matcher is None:
1231 # Should exactly be None
1232 return {} if node is None else None
1233
1234 if isinstance(matcher, str):
1235 # Should exactly match matcher text
1236 return {} if node == matcher else None
1237
1238 if isinstance(matcher, bool):
1239 # Should exactly match matcher bool
1240 return {} if node is matcher else None
1241
1242 if isinstance(node, collections.abc.Sequence):
1243 # Given we've generated the types for matchers based on LibCST, we know that
1244 # this is true unless the node is badly constructed and types were ignored.
1245 node = cast(Sequence[Union[MaybeSentinel, libcst.CSTNode]], node)
1246
1247 if isinstance(matcher, OneOf):
1248 # We should compare against each of the sequences in the OneOf
1249 for m in matcher.options:
1250 if isinstance(m, collections.abc.Sequence):
1251 # Should match the sequence of requested nodes
1252 result = _sequence_matches(node, m, metadata_lookup)
1253 if result.sequence_capture is not None:
1254 return result.sequence_capture
1255 elif isinstance(m, MatchIfTrue):
1256 # TODO: return captures
1257 return {} if m.func(node) else None
1258 elif isinstance(matcher, AllOf):
1259 # We should compare against each of the sequences in the AllOf
1260 all_captures = {}
1261 for m in matcher.options:
1262 if isinstance(m, collections.abc.Sequence):
1263 # Should match the sequence of requested nodes
1264 result = _sequence_matches(node, m, metadata_lookup)
1265 if result.sequence_capture is None:
1266 return None
1267 all_captures = {**all_captures, **result.sequence_capture}
1268 else:
1269 # The value in the AllOf wasn't a sequence, it can't match.
1270 return None
1271 # We passed the checks above for each node, so we passed.
1272 return all_captures
1273 elif isinstance(matcher, collections.abc.Sequence):
1274 # We should assume that this matcher is a sequence to compare. Given
1275 # the way we generate match classes, this should be true unless the
1276 # match is badly constructed and types were ignored.
1277 return _sequence_matches(
1278 node,
1279 cast(
1280 Sequence[
1281 Union[
1282 BaseMatcherNode,
1283 _BaseWildcardNode,
1284 MatchIfTrue[libcst.CSTNode],
1285 DoNotCareSentinel,
1286 ]
1287 ],
1288 matcher,
1289 ),
1290 metadata_lookup,
1291 ).sequence_capture
1292
1293 # We exhausted our possibilities, there's no match
1294 return None
1295
1296 # Base case, should match node via matcher. We know the type of node is
1297 # correct here because we generate matchers directly off of LibCST nodes,
1298 # so the only way it is wrong is if the node was badly constructed and
1299 # types were ignored.
1300 return _matches(
1301 cast(Union[MaybeSentinel, libcst.CSTNode], node),
1302 # pyre-fixme[24]: Generic type `MatchIfTrue` expects 1 type parameter.
1303 cast(Union[BaseMatcherNode, MatchIfTrue, _BaseMetadataMatcher], matcher),
1304 metadata_lookup,
1305 )
1306
1307
1308def _metadata_matches( # noqa: C901
1309 node: libcst.CSTNode,
1310 metadata: Union[
1311 _BaseMetadataMatcher,
1312 AllOf[_BaseMetadataMatcher],
1313 OneOf[_BaseMetadataMatcher],
1314 _InverseOf[_BaseMetadataMatcher],
1315 _ExtractMatchingNode[_BaseMetadataMatcher],
1316 ],
1317 metadata_lookup: Callable[[meta.ProviderT, libcst.CSTNode], object],
1318) -> Optional[Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]]:
1319 if isinstance(metadata, OneOf):
1320 for metadata in metadata.options:
1321 metadata_capture = _metadata_matches(node, metadata, metadata_lookup)
1322 if metadata_capture is not None:
1323 return metadata_capture
1324 return None
1325 elif isinstance(metadata, AllOf):
1326 all_captures = {}
1327 for metadata in metadata.options:
1328 metadata_capture = _metadata_matches(node, metadata, metadata_lookup)
1329 if metadata_capture is None:
1330 return None
1331 all_captures = {**all_captures, **metadata_capture}
1332 # We passed the above checks, so we pass the matcher.
1333 return all_captures
1334 elif isinstance(metadata, _InverseOf):
1335 return (
1336 {}
1337 if _metadata_matches(node, metadata.matcher, metadata_lookup) is None
1338 else None
1339 )
1340 elif isinstance(metadata, _ExtractMatchingNode):
1341 metadata_capture = _metadata_matches(node, metadata.matcher, metadata_lookup)
1342 if metadata_capture is not None:
1343 return {
1344 # Our own match capture comes last, since its higher in the tree
1345 # so we want to override any child match captures by the same name.
1346 **metadata_capture,
1347 metadata.name: node,
1348 }
1349 return None
1350 elif isinstance(metadata, MatchMetadataIfTrue):
1351 actual_value = metadata_lookup(metadata.key, node)
1352 if actual_value is _METADATA_MISSING_SENTINEL:
1353 return None
1354 return {} if metadata.func(actual_value) else None
1355 elif isinstance(metadata, MatchMetadata):
1356 actual_value = metadata_lookup(metadata.key, node)
1357 if actual_value is _METADATA_MISSING_SENTINEL:
1358 return None
1359 return {} if actual_value == metadata.value else None
1360 else:
1361 raise CSTLogicError("Logic error!")
1362
1363
1364def _node_matches( # noqa: C901
1365 node: libcst.CSTNode,
1366 matcher: Union[
1367 BaseMatcherNode,
1368 MatchIfTrue[libcst.CSTNode],
1369 _BaseMetadataMatcher,
1370 _InverseOf[
1371 Union[
1372 BaseMatcherNode,
1373 MatchIfTrue[libcst.CSTNode],
1374 _BaseMetadataMatcher,
1375 ]
1376 ],
1377 _ExtractMatchingNode[
1378 Union[
1379 BaseMatcherNode,
1380 MatchIfTrue[libcst.CSTNode],
1381 _BaseMetadataMatcher,
1382 ]
1383 ],
1384 ],
1385 metadata_lookup: Callable[[meta.ProviderT, libcst.CSTNode], object],
1386) -> Optional[Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]]:
1387 # If this is a _InverseOf, then invert the result.
1388 if isinstance(matcher, _InverseOf):
1389 return (
1390 {}
1391 if _node_matches(node, matcher.matcher, metadata_lookup) is None
1392 else None
1393 )
1394
1395 # If this is an _ExtractMatchingNode, grab the resulting call and pass the check
1396 # forward.
1397 if isinstance(matcher, _ExtractMatchingNode):
1398 node_capture = _node_matches(node, matcher.matcher, metadata_lookup)
1399 if node_capture is not None:
1400 return {
1401 # We come last here since we're further up the tree, so we want to
1402 # override any identically named child match nodes.
1403 **node_capture,
1404 matcher.name: node,
1405 }
1406 return None
1407
1408 # Now, check if this is a lambda matcher.
1409 if isinstance(matcher, MatchIfTrue):
1410 return {} if matcher.func(node) else None
1411
1412 if isinstance(matcher, (MatchMetadata, MatchMetadataIfTrue)):
1413 return _metadata_matches(node, matcher, metadata_lookup)
1414
1415 # Now, check that the node and matcher classes are the same.
1416 if node.__class__.__name__ != matcher.__class__.__name__:
1417 return None
1418
1419 # Now, check that the children match for each attribute.
1420 all_captures = {}
1421 for field in fields(matcher):
1422 if field.name == "_metadata":
1423 # We don't care about this field, its a dataclasses implementation detail.
1424 continue
1425 elif field.name == "metadata":
1426 # Special field we respect for matching metadata on a particular node.
1427 desired = getattr(matcher, field.name)
1428 if isinstance(desired, DoNotCareSentinel):
1429 # We don't care about this
1430 continue
1431 metadata_capture = _metadata_matches(node, desired, metadata_lookup)
1432 if metadata_capture is None:
1433 return None
1434 all_captures = {**all_captures, **metadata_capture}
1435 else:
1436 desired = getattr(matcher, field.name)
1437 actual = getattr(node, field.name)
1438 attribute_capture = _attribute_matches(actual, desired, metadata_lookup)
1439 if attribute_capture is None:
1440 return None
1441 all_captures = {**all_captures, **attribute_capture}
1442
1443 # We didn't find a non-match in the above loop, so it matches!
1444 return all_captures
1445
1446
1447def _matches(
1448 node: Union[MaybeSentinel, libcst.CSTNode],
1449 matcher: Union[
1450 BaseMatcherNode,
1451 MatchIfTrue[libcst.CSTNode],
1452 _BaseMetadataMatcher,
1453 _InverseOf[
1454 Union[
1455 BaseMatcherNode,
1456 MatchIfTrue[libcst.CSTNode],
1457 _BaseMetadataMatcher,
1458 ]
1459 ],
1460 _ExtractMatchingNode[
1461 Union[
1462 BaseMatcherNode,
1463 MatchIfTrue[libcst.CSTNode],
1464 _BaseMetadataMatcher,
1465 ]
1466 ],
1467 ],
1468 metadata_lookup: Callable[[meta.ProviderT, libcst.CSTNode], object],
1469) -> Optional[Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]]:
1470 if isinstance(node, MaybeSentinel):
1471 # We can't possibly match on a maybe sentinel, so it only matches if
1472 # the matcher we have is a _InverseOf.
1473 return {} if isinstance(matcher, _InverseOf) else None
1474
1475 # Now, evaluate the matcher node itself.
1476 if isinstance(matcher, (OneOf, TypeOf)):
1477 for matcher in matcher.options:
1478 node_capture = _node_matches(node, matcher, metadata_lookup)
1479 if node_capture is not None:
1480 return node_capture
1481 return None
1482 elif isinstance(matcher, AllOf):
1483 all_captures = {}
1484 for matcher in matcher.options:
1485 node_capture = _node_matches(node, matcher, metadata_lookup)
1486 if node_capture is None:
1487 return None
1488 all_captures = {**all_captures, **node_capture}
1489 return all_captures
1490 else:
1491 return _node_matches(node, matcher, metadata_lookup)
1492
1493
1494def _construct_metadata_fetcher_null() -> (
1495 Callable[[meta.ProviderT, libcst.CSTNode], object]
1496):
1497 def _fetch(provider: meta.ProviderT, node: libcst.CSTNode) -> NoReturn:
1498 raise LookupError(
1499 f"{provider.__name__} is not resolved; did you forget a MetadataWrapper?"
1500 )
1501
1502 return _fetch
1503
1504
1505def _construct_metadata_fetcher_dependent(
1506 dependent_class: libcst.MetadataDependent,
1507) -> Callable[[meta.ProviderT, libcst.CSTNode], object]:
1508 def _fetch(provider: meta.ProviderT, node: libcst.CSTNode) -> object:
1509 return dependent_class.get_metadata(provider, node, _METADATA_MISSING_SENTINEL)
1510
1511 return _fetch
1512
1513
1514def _construct_metadata_fetcher_wrapper(
1515 wrapper: libcst.MetadataWrapper,
1516) -> Callable[[meta.ProviderT, libcst.CSTNode], object]:
1517 metadata: Dict[meta.ProviderT, Mapping[libcst.CSTNode, object]] = {}
1518
1519 def _fetch(provider: meta.ProviderT, node: libcst.CSTNode) -> object:
1520 if provider not in metadata:
1521 metadata[provider] = wrapper.resolve(provider)
1522
1523 node_metadata = metadata[provider].get(node, _METADATA_MISSING_SENTINEL)
1524 if isinstance(node_metadata, LazyValue):
1525 node_metadata = node_metadata()
1526
1527 return node_metadata
1528
1529 return _fetch
1530
1531
1532def extract(
1533 node: Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode],
1534 matcher: BaseMatcherNode,
1535 *,
1536 metadata_resolver: Optional[
1537 Union[libcst.MetadataDependent, libcst.MetadataWrapper]
1538 ] = None,
1539) -> Optional[Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]]:
1540 """
1541 Given an arbitrary node from a LibCST tree, and an arbitrary matcher, returns
1542 a dictionary of extracted children of the tree if the node matches the shape defined
1543 by the matcher. Note that the node can also be a :class:`~libcst.RemovalSentinel` or
1544 a :class:`~libcst.MaybeSentinel` in order to use extract directly on transform results
1545 and node attributes. In these cases, :func:`extract` will always return ``None``.
1546
1547 If the node matches the shape defined by the matcher, the return will be a dictionary
1548 whose keys are defined by the :func:`SaveMatchedNode` name parameter, and the values
1549 will be the node or sequence that was present at that location in the shape defined
1550 by the matcher. In the case of multiple :func:`SaveMatchedNode` matches with the
1551 same name, parent nodes will take prioirity over child nodes, and nodes later in
1552 sequences will take priority over nodes earlier in sequences.
1553
1554 The matcher can be any concrete matcher that subclasses from :class:`BaseMatcherNode`,
1555 or a :class:`OneOf`/:class:`AllOf` special matcher. It cannot be a
1556 :class:`MatchIfTrue` or a :func:`DoesNotMatch` matcher since these are redundant.
1557 It cannot be a :class:`AtLeastN` or :class:`AtMostN` matcher because these types are
1558 wildcards which can only be used inside sequences.
1559 """
1560 if isinstance(node, RemovalSentinel):
1561 # We can't possibly match on a removal sentinel, so it doesn't match.
1562 return None
1563 if isinstance(matcher, (AtLeastN, AtMostN, MatchIfTrue, _BaseMetadataMatcher)):
1564 # We can't match this, since these matchers are forbidden at top level.
1565 # These are not subclasses of BaseMatcherNode, but in the case that the
1566 # user is not using type checking, this should still behave correctly.
1567 return None
1568
1569 if metadata_resolver is None:
1570 fetcher = _construct_metadata_fetcher_null()
1571 elif isinstance(metadata_resolver, libcst.MetadataWrapper):
1572 fetcher = _construct_metadata_fetcher_wrapper(metadata_resolver)
1573 else:
1574 fetcher = _construct_metadata_fetcher_dependent(metadata_resolver)
1575
1576 return _matches(node, matcher, fetcher)
1577
1578
1579def matches(
1580 node: Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode],
1581 matcher: BaseMatcherNode,
1582 *,
1583 metadata_resolver: Optional[
1584 Union[libcst.MetadataDependent, libcst.MetadataWrapper]
1585 ] = None,
1586) -> bool:
1587 """
1588 Given an arbitrary node from a LibCST tree, and an arbitrary matcher, returns
1589 ``True`` if the node matches the shape defined by the matcher. Note that the node
1590 can also be a :class:`~libcst.RemovalSentinel` or a :class:`~libcst.MaybeSentinel`
1591 in order to use matches directly on transform results and node attributes. In these
1592 cases, :func:`matches` will always return ``False``.
1593
1594 The matcher can be any concrete matcher that subclasses from :class:`BaseMatcherNode`,
1595 or a :class:`OneOf`/:class:`AllOf` special matcher. It cannot be a
1596 :class:`MatchIfTrue` or a :func:`DoesNotMatch` matcher since these are redundant.
1597 It cannot be a :class:`AtLeastN` or :class:`AtMostN` matcher because these types
1598 are wildcards which can only be used inside sequences.
1599 """
1600 return extract(node, matcher, metadata_resolver=metadata_resolver) is not None
1601
1602
1603class _FindAllVisitor(libcst.CSTVisitor):
1604 def __init__(
1605 self,
1606 matcher: Union[
1607 BaseMatcherNode,
1608 MatchIfTrue[libcst.CSTNode],
1609 _BaseMetadataMatcher,
1610 _InverseOf[
1611 Union[
1612 BaseMatcherNode,
1613 MatchIfTrue[libcst.CSTNode],
1614 _BaseMetadataMatcher,
1615 ]
1616 ],
1617 ],
1618 metadata_lookup: Callable[[meta.ProviderT, libcst.CSTNode], object],
1619 ) -> None:
1620 self.matcher = matcher
1621 self.metadata_lookup = metadata_lookup
1622 self.found_nodes: List[libcst.CSTNode] = []
1623 self.extracted_nodes: List[
1624 Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]
1625 ] = []
1626
1627 def on_visit(self, node: libcst.CSTNode) -> bool:
1628 match = _matches(node, self.matcher, self.metadata_lookup)
1629 if match is not None:
1630 self.found_nodes.append(node)
1631 self.extracted_nodes.append(match)
1632 return True
1633
1634
1635def _find_or_extract_all(
1636 tree: Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode, meta.MetadataWrapper],
1637 matcher: Union[
1638 BaseMatcherNode,
1639 MatchIfTrue[libcst.CSTNode],
1640 _BaseMetadataMatcher,
1641 # The inverse clause is left off of the public functions `findall` and
1642 # `extractall` because we play a dirty trick. We lie to the typechecker
1643 # that `DoesNotMatch` returns identity, so the public functions don't
1644 # need to be aware of inverses. If we could represent predicate logic
1645 # in python types we could get away with this, but that's not the state
1646 # of things right now.
1647 _InverseOf[
1648 Union[
1649 BaseMatcherNode,
1650 MatchIfTrue[libcst.CSTNode],
1651 _BaseMetadataMatcher,
1652 ]
1653 ],
1654 ],
1655 *,
1656 metadata_resolver: Optional[
1657 Union[libcst.MetadataDependent, libcst.MetadataWrapper]
1658 ] = None,
1659) -> Tuple[
1660 Sequence[libcst.CSTNode],
1661 Sequence[Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]],
1662]:
1663 if isinstance(tree, (RemovalSentinel, MaybeSentinel)):
1664 # We can't possibly match on a removal sentinel, so it doesn't match.
1665 return [], []
1666 if isinstance(matcher, (AtLeastN, AtMostN)):
1667 # We can't match this, since these matchers are forbidden at top level.
1668 # These are not subclasses of BaseMatcherNode, but in the case that the
1669 # user is not using type checking, this should still behave correctly.
1670 return [], []
1671
1672 if isinstance(tree, meta.MetadataWrapper) and metadata_resolver is None:
1673 # Provide a convenience for calling findall directly on a MetadataWrapper.
1674 metadata_resolver = tree
1675
1676 if metadata_resolver is None:
1677 fetcher = _construct_metadata_fetcher_null()
1678 elif isinstance(metadata_resolver, libcst.MetadataWrapper):
1679 fetcher = _construct_metadata_fetcher_wrapper(metadata_resolver)
1680 else:
1681 fetcher = _construct_metadata_fetcher_dependent(metadata_resolver)
1682
1683 finder = _FindAllVisitor(matcher, fetcher)
1684 tree.visit(finder)
1685 return finder.found_nodes, finder.extracted_nodes
1686
1687
1688def findall(
1689 tree: Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode, meta.MetadataWrapper],
1690 matcher: Union[BaseMatcherNode, MatchIfTrue[libcst.CSTNode], _BaseMetadataMatcher],
1691 *,
1692 metadata_resolver: Optional[
1693 Union[libcst.MetadataDependent, libcst.MetadataWrapper]
1694 ] = None,
1695) -> Sequence[libcst.CSTNode]:
1696 """
1697 Given an arbitrary node from a LibCST tree and an arbitrary matcher, iterates
1698 over that node and all children returning a sequence of all child nodes that
1699 match the given matcher. Note that the tree can also be a
1700 :class:`~libcst.RemovalSentinel` or a :class:`~libcst.MaybeSentinel`
1701 in order to use findall directly on transform results and node attributes. In these
1702 cases, :func:`findall` will always return an empty sequence. Note also that
1703 instead of a LibCST tree, you can instead pass in a
1704 :class:`~libcst.metadata.MetadataWrapper`. This mirrors the fact that you can
1705 call ``visit`` on a :class:`~libcst.metadata.MetadataWrapper` in order to iterate
1706 over it with a transform. If you provide a wrapper for the tree and do not set
1707 the ``metadata_resolver`` parameter specifically, it will automatically be set
1708 to the wrapper for you.
1709
1710 The matcher can be any concrete matcher that subclasses from :class:`BaseMatcherNode`,
1711 or a :class:`OneOf`/:class:`AllOf` special matcher. Unlike :func:`matches`, it can
1712 also be a :class:`MatchIfTrue` or :func:`DoesNotMatch` matcher, since we are
1713 traversing the tree looking for matches. It cannot be a :class:`AtLeastN` or
1714 :class:`AtMostN` matcher because these types are wildcards which can only be used
1715 inside sequences.
1716 """
1717 nodes, _ = _find_or_extract_all(tree, matcher, metadata_resolver=metadata_resolver)
1718 return nodes
1719
1720
1721def extractall(
1722 tree: Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode, meta.MetadataWrapper],
1723 matcher: Union[BaseMatcherNode, MatchIfTrue[libcst.CSTNode], _BaseMetadataMatcher],
1724 *,
1725 metadata_resolver: Optional[
1726 Union[libcst.MetadataDependent, libcst.MetadataWrapper]
1727 ] = None,
1728) -> Sequence[Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]]:
1729 """
1730 Given an arbitrary node from a LibCST tree and an arbitrary matcher, iterates
1731 over that node and all children returning a sequence of dictionaries representing
1732 the saved and extracted children specified by :func:`SaveMatchedNode` for each
1733 match found in the tree. This is analogous to running a :func:`findall` over a
1734 tree, then running :func:`extract` with the same matcher over each of the returned
1735 nodes. Note that the tree can also be a :class:`~libcst.RemovalSentinel` or a
1736 :class:`~libcst.MaybeSentinel` in order to use extractall directly on transform
1737 results and node attributes. In these cases, :func:`extractall` will always
1738 return an empty sequence. Note also that instead of a LibCST tree, you can
1739 instead pass in a :class:`~libcst.metadata.MetadataWrapper`. This mirrors the
1740 fact that you can call ``visit`` on a :class:`~libcst.metadata.MetadataWrapper`
1741 in order to iterate over it with a transform. If you provide a wrapper for the
1742 tree and do not set the ``metadata_resolver`` parameter specifically, it will
1743 automatically be set to the wrapper for you.
1744
1745 The matcher can be any concrete matcher that subclasses from :class:`BaseMatcherNode`,
1746 or a :class:`OneOf`/:class:`AllOf` special matcher. Unlike :func:`matches`, it can
1747 also be a :class:`MatchIfTrue` or :func:`DoesNotMatch` matcher, since we are
1748 traversing the tree looking for matches. It cannot be a :class:`AtLeastN` or
1749 :class:`AtMostN` matcher because these types are wildcards which can only be usedi
1750 inside sequences.
1751 """
1752 _, extractions = _find_or_extract_all(
1753 tree, matcher, metadata_resolver=metadata_resolver
1754 )
1755 return extractions
1756
1757
1758class _ReplaceTransformer(libcst.CSTTransformer):
1759 def __init__(
1760 self,
1761 matcher: Union[
1762 BaseMatcherNode,
1763 MatchIfTrue[libcst.CSTNode],
1764 _BaseMetadataMatcher,
1765 _InverseOf[
1766 Union[
1767 BaseMatcherNode,
1768 MatchIfTrue[libcst.CSTNode],
1769 _BaseMetadataMatcher,
1770 ]
1771 ],
1772 ],
1773 metadata_lookup: Callable[[meta.ProviderT, libcst.CSTNode], object],
1774 replacement: Union[
1775 MaybeSentinel,
1776 RemovalSentinel,
1777 libcst.CSTNode,
1778 Callable[
1779 [
1780 libcst.CSTNode,
1781 Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]],
1782 ],
1783 Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode],
1784 ],
1785 ],
1786 ) -> None:
1787 self.matcher = matcher
1788 self.metadata_lookup = metadata_lookup
1789 self.replacement: Callable[
1790 [
1791 libcst.CSTNode,
1792 Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]],
1793 ],
1794 Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode],
1795 ]
1796
1797 if inspect.isfunction(replacement):
1798 self.replacement = replacement
1799 elif isinstance(replacement, (MaybeSentinel, RemovalSentinel)):
1800 self.replacement = lambda node, matches: replacement
1801 else:
1802 # pyre-ignore We know this is a CSTNode.
1803 self.replacement = lambda node, matches: replacement.deep_clone()
1804 # We run into a really weird problem here, where we need to run the match
1805 # and extract step on the original node in order for metadata to work.
1806 # However, if we do that, then using things like `deep_replace` will fail
1807 # since any extracted nodes are the originals, not the updates and LibCST
1808 # does replacement by identity for safety reasons. If we try to run the
1809 # match and extract step on the updated node (or twice, once for the match
1810 # and once for the extract), it will fail to extract if any metadata-based
1811 # matchers are used. So, we try to compromise with the best of both worlds.
1812 # We track all node updates, and when we send the extracted nodes to the
1813 # replacement callable, we look up the original nodes and replace them with
1814 # updated nodes. In the case that an update made the node no-longer exist,
1815 # we act as if there was not a match (because in reality, there would not
1816 # have been if we had run the matcher on the update).
1817 self.node_lut: Dict[libcst.CSTNode, libcst.CSTNode] = {}
1818
1819 def _node_translate(
1820 self, node_or_sequence: Union[libcst.CSTNode, Sequence[libcst.CSTNode]]
1821 ) -> Union[libcst.CSTNode, Sequence[libcst.CSTNode]]:
1822 if isinstance(node_or_sequence, Sequence):
1823 return tuple(self.node_lut[node] for node in node_or_sequence)
1824 else:
1825 return self.node_lut[node_or_sequence]
1826
1827 def _extraction_translate(
1828 self, extracted: Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]
1829 ) -> Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]]:
1830 return {key: self._node_translate(val) for key, val in extracted.items()}
1831
1832 def on_leave(
1833 self, original_node: libcst.CSTNode, updated_node: libcst.CSTNode
1834 ) -> Union[libcst.CSTNode, MaybeSentinel, RemovalSentinel]:
1835 # Track original to updated node mapping for this node.
1836 self.node_lut[original_node] = updated_node
1837
1838 # This gets complicated. We need to do the match on the original node,
1839 # but we want to do the extraction on the updated node. This is so
1840 # metadata works properly in matchers. So, if we get a match, we fix
1841 # up the nodes in the match and return that to the replacement lambda.
1842 extracted = _matches(original_node, self.matcher, self.metadata_lookup)
1843 if extracted is not None:
1844 try:
1845 # Attempt to do a translation from original to updated node.
1846 extracted = self._extraction_translate(extracted)
1847 except KeyError:
1848 # One of the nodes we looked up doesn't exist anymore, this
1849 # is no longer a match. This can happen if a child node was
1850 # modified, making this original match not applicable anymore.
1851 extracted = None
1852 if extracted is not None:
1853 # We're replacing this node entirely, so don't save the original
1854 # updated node. We don't want this to be part of a parent match
1855 # since we can't guarantee that the update matches anymore.
1856 del self.node_lut[original_node]
1857 return self.replacement(updated_node, extracted)
1858 return updated_node
1859
1860
1861def replace(
1862 tree: Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode, meta.MetadataWrapper],
1863 matcher: Union[BaseMatcherNode, MatchIfTrue[libcst.CSTNode], _BaseMetadataMatcher],
1864 replacement: Union[
1865 MaybeSentinel,
1866 RemovalSentinel,
1867 libcst.CSTNode,
1868 Callable[
1869 [
1870 libcst.CSTNode,
1871 Dict[str, Union[libcst.CSTNode, Sequence[libcst.CSTNode]]],
1872 ],
1873 Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode],
1874 ],
1875 ],
1876 *,
1877 metadata_resolver: Optional[
1878 Union[libcst.MetadataDependent, libcst.MetadataWrapper]
1879 ] = None,
1880) -> Union[MaybeSentinel, RemovalSentinel, libcst.CSTNode]:
1881 """
1882 Given an arbitrary node from a LibCST tree and an arbitrary matcher, iterates
1883 over that node and all children and replaces each node that matches the supplied
1884 matcher with a supplied replacement. Note that the replacement can either be a
1885 valid node type, or a callable which takes the matched node and a dictionary of
1886 any extracted child values and returns a valid node type. If you provide a
1887 valid LibCST node type, :func:`replace` will replace every node that matches
1888 the supplied matcher with the replacement node. If you provide a callable,
1889 :func:`replace` will run :func:`extract` over all matched nodes and call the
1890 callable with both the node that should be replaced and the dictionary returned
1891 by :func:`extract`. Under all circumstances a new tree is returned.
1892 :func:`extract` should be viewed as a short-cut to writing a transform which
1893 also returns a new tree even when no changes are applied.
1894
1895 Note that the tree can also be a :class:`~libcst.RemovalSentinel` or a
1896 :class:`~libcst.MaybeSentinel` in order to use replace directly on transform
1897 results and node attributes. In these cases, :func:`replace` will return the
1898 same :class:`~libcst.RemovalSentinel` or :class:`~libcst.MaybeSentinel`.
1899 Note also that instead of a LibCST tree, you can instead pass in a
1900 :class:`~libcst.metadata.MetadataWrapper`. This mirrors the fact that you can
1901 call ``visit`` on a :class:`~libcst.metadata.MetadataWrapper` in order to
1902 iterate over it with a transform. If you provide a wrapper for the tree and
1903 do not set the ``metadata_resolver`` parameter specifically, it will
1904 automatically be set to the wrapper for you.
1905
1906 The matcher can be any concrete matcher that subclasses from :class:`BaseMatcherNode`,
1907 or a :class:`OneOf`/:class:`AllOf` special matcher. Unlike :func:`matches`, it can
1908 also be a :class:`MatchIfTrue` or :func:`DoesNotMatch` matcher, since we are
1909 traversing the tree looking for matches. It cannot be a :class:`AtLeastN` or
1910 :class:`AtMostN` matcher because these types are wildcards which can only be usedi
1911 inside sequences.
1912 """
1913 if isinstance(tree, (RemovalSentinel, MaybeSentinel)):
1914 # We can't do any replacements on this, so return the tree exactly.
1915 return tree
1916 if isinstance(matcher, (AtLeastN, AtMostN)):
1917 # We can't match this, since these matchers are forbidden at top level.
1918 # These are not subclasses of BaseMatcherNode, but in the case that the
1919 # user is not using type checking, this should still behave correctly.
1920 if isinstance(tree, libcst.CSTNode):
1921 return tree.deep_clone()
1922 elif isinstance(tree, meta.MetadataWrapper):
1923 return tree.module.deep_clone()
1924 else:
1925 raise CSTLogicError("Logic error!")
1926
1927 if isinstance(tree, meta.MetadataWrapper) and metadata_resolver is None:
1928 # Provide a convenience for calling replace directly on a MetadataWrapper.
1929 metadata_resolver = tree
1930
1931 if metadata_resolver is None:
1932 fetcher = _construct_metadata_fetcher_null()
1933 elif isinstance(metadata_resolver, libcst.MetadataWrapper):
1934 fetcher = _construct_metadata_fetcher_wrapper(metadata_resolver)
1935 else:
1936 fetcher = _construct_metadata_fetcher_dependent(metadata_resolver)
1937
1938 replacer = _ReplaceTransformer(matcher, fetcher, replacement)
1939 new_tree = tree.visit(replacer)
1940 if isinstance(new_tree, FlattenSentinel):
1941 # The above transform never returns FlattenSentinel, so this isn't possible
1942 raise CSTLogicError("Logic error, cannot get a FlattenSentinel here!")
1943 return new_tree