visitors.py 35 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164
  1. # sql/visitors.py
  2. # Copyright (C) 2005-2025 the SQLAlchemy authors and contributors
  3. # <see AUTHORS file>
  4. #
  5. # This module is part of SQLAlchemy and is released under
  6. # the MIT License: https://www.opensource.org/licenses/mit-license.php
  7. """Visitor/traversal interface and library functions."""
  8. from __future__ import annotations
  9. from collections import deque
  10. from enum import Enum
  11. import itertools
  12. import operator
  13. import typing
  14. from typing import Any
  15. from typing import Callable
  16. from typing import cast
  17. from typing import ClassVar
  18. from typing import Dict
  19. from typing import Iterable
  20. from typing import Iterator
  21. from typing import List
  22. from typing import Mapping
  23. from typing import Optional
  24. from typing import overload
  25. from typing import Tuple
  26. from typing import Type
  27. from typing import TYPE_CHECKING
  28. from typing import TypeVar
  29. from typing import Union
  30. from .. import exc
  31. from .. import util
  32. from ..util import langhelpers
  33. from ..util._has_cy import HAS_CYEXTENSION
  34. from ..util.typing import Literal
  35. from ..util.typing import Protocol
  36. from ..util.typing import Self
  37. if TYPE_CHECKING:
  38. from .annotation import _AnnotationDict
  39. from .elements import ColumnElement
  40. if typing.TYPE_CHECKING or not HAS_CYEXTENSION:
  41. from ._py_util import prefix_anon_map as prefix_anon_map
  42. from ._py_util import cache_anon_map as anon_map
  43. else:
  44. from sqlalchemy.cyextension.util import ( # noqa: F401,E501
  45. prefix_anon_map as prefix_anon_map,
  46. )
  47. from sqlalchemy.cyextension.util import ( # noqa: F401,E501
  48. cache_anon_map as anon_map,
  49. )
  50. __all__ = [
  51. "iterate",
  52. "traverse_using",
  53. "traverse",
  54. "cloned_traverse",
  55. "replacement_traverse",
  56. "Visitable",
  57. "ExternalTraversal",
  58. "InternalTraversal",
  59. "anon_map",
  60. ]
  61. class _CompilerDispatchType(Protocol):
  62. def __call__(_self, self: Visitable, visitor: Any, **kw: Any) -> Any: ...
  63. class Visitable:
  64. """Base class for visitable objects.
  65. :class:`.Visitable` is used to implement the SQL compiler dispatch
  66. functions. Other forms of traversal such as for cache key generation
  67. are implemented separately using the :class:`.HasTraverseInternals`
  68. interface.
  69. .. versionchanged:: 2.0 The :class:`.Visitable` class was named
  70. :class:`.Traversible` in the 1.4 series; the name is changed back
  71. to :class:`.Visitable` in 2.0 which is what it was prior to 1.4.
  72. Both names remain importable in both 1.4 and 2.0 versions.
  73. """
  74. __slots__ = ()
  75. __visit_name__: str
  76. _original_compiler_dispatch: _CompilerDispatchType
  77. if typing.TYPE_CHECKING:
  78. def _compiler_dispatch(self, visitor: Any, **kw: Any) -> str: ...
  79. def __init_subclass__(cls) -> None:
  80. if "__visit_name__" in cls.__dict__:
  81. cls._generate_compiler_dispatch()
  82. super().__init_subclass__()
  83. @classmethod
  84. def _generate_compiler_dispatch(cls) -> None:
  85. visit_name = cls.__visit_name__
  86. if "_compiler_dispatch" in cls.__dict__:
  87. # class has a fixed _compiler_dispatch() method.
  88. # copy it to "original" so that we can get it back if
  89. # sqlalchemy.ext.compiles overrides it.
  90. cls._original_compiler_dispatch = cls._compiler_dispatch
  91. return
  92. if not isinstance(visit_name, str):
  93. raise exc.InvalidRequestError(
  94. f"__visit_name__ on class {cls.__name__} must be a string "
  95. "at the class level"
  96. )
  97. name = "visit_%s" % visit_name
  98. getter = operator.attrgetter(name)
  99. def _compiler_dispatch(
  100. self: Visitable, visitor: Any, **kw: Any
  101. ) -> str:
  102. """Look for an attribute named "visit_<visit_name>" on the
  103. visitor, and call it with the same kw params.
  104. """
  105. try:
  106. meth = getter(visitor)
  107. except AttributeError as err:
  108. return visitor.visit_unsupported_compilation(self, err, **kw) # type: ignore # noqa: E501
  109. else:
  110. return meth(self, **kw) # type: ignore # noqa: E501
  111. cls._compiler_dispatch = ( # type: ignore
  112. cls._original_compiler_dispatch
  113. ) = _compiler_dispatch
  114. def __class_getitem__(cls, key: Any) -> Any:
  115. # allow generic classes in py3.9+
  116. return cls
  117. class InternalTraversal(Enum):
  118. r"""Defines visitor symbols used for internal traversal.
  119. The :class:`.InternalTraversal` class is used in two ways. One is that
  120. it can serve as the superclass for an object that implements the
  121. various visit methods of the class. The other is that the symbols
  122. themselves of :class:`.InternalTraversal` are used within
  123. the ``_traverse_internals`` collection. Such as, the :class:`.Case`
  124. object defines ``_traverse_internals`` as ::
  125. class Case(ColumnElement[_T]):
  126. _traverse_internals = [
  127. ("value", InternalTraversal.dp_clauseelement),
  128. ("whens", InternalTraversal.dp_clauseelement_tuples),
  129. ("else_", InternalTraversal.dp_clauseelement),
  130. ]
  131. Above, the :class:`.Case` class indicates its internal state as the
  132. attributes named ``value``, ``whens``, and ``else_``. They each
  133. link to an :class:`.InternalTraversal` method which indicates the type
  134. of datastructure to which each attribute refers.
  135. Using the ``_traverse_internals`` structure, objects of type
  136. :class:`.InternalTraversible` will have the following methods automatically
  137. implemented:
  138. * :meth:`.HasTraverseInternals.get_children`
  139. * :meth:`.HasTraverseInternals._copy_internals`
  140. * :meth:`.HasCacheKey._gen_cache_key`
  141. Subclasses can also implement these methods directly, particularly for the
  142. :meth:`.HasTraverseInternals._copy_internals` method, when special steps
  143. are needed.
  144. .. versionadded:: 1.4
  145. """
  146. dp_has_cache_key = "HC"
  147. """Visit a :class:`.HasCacheKey` object."""
  148. dp_has_cache_key_list = "HL"
  149. """Visit a list of :class:`.HasCacheKey` objects."""
  150. dp_clauseelement = "CE"
  151. """Visit a :class:`_expression.ClauseElement` object."""
  152. dp_fromclause_canonical_column_collection = "FC"
  153. """Visit a :class:`_expression.FromClause` object in the context of the
  154. ``columns`` attribute.
  155. The column collection is "canonical", meaning it is the originally
  156. defined location of the :class:`.ColumnClause` objects. Right now
  157. this means that the object being visited is a
  158. :class:`_expression.TableClause`
  159. or :class:`_schema.Table` object only.
  160. """
  161. dp_clauseelement_tuples = "CTS"
  162. """Visit a list of tuples which contain :class:`_expression.ClauseElement`
  163. objects.
  164. """
  165. dp_clauseelement_list = "CL"
  166. """Visit a list of :class:`_expression.ClauseElement` objects.
  167. """
  168. dp_clauseelement_tuple = "CT"
  169. """Visit a tuple of :class:`_expression.ClauseElement` objects.
  170. """
  171. dp_executable_options = "EO"
  172. dp_with_context_options = "WC"
  173. dp_fromclause_ordered_set = "CO"
  174. """Visit an ordered set of :class:`_expression.FromClause` objects. """
  175. dp_string = "S"
  176. """Visit a plain string value.
  177. Examples include table and column names, bound parameter keys, special
  178. keywords such as "UNION", "UNION ALL".
  179. The string value is considered to be significant for cache key
  180. generation.
  181. """
  182. dp_string_list = "SL"
  183. """Visit a list of strings."""
  184. dp_anon_name = "AN"
  185. """Visit a potentially "anonymized" string value.
  186. The string value is considered to be significant for cache key
  187. generation.
  188. """
  189. dp_boolean = "B"
  190. """Visit a boolean value.
  191. The boolean value is considered to be significant for cache key
  192. generation.
  193. """
  194. dp_operator = "O"
  195. """Visit an operator.
  196. The operator is a function from the :mod:`sqlalchemy.sql.operators`
  197. module.
  198. The operator value is considered to be significant for cache key
  199. generation.
  200. """
  201. dp_type = "T"
  202. """Visit a :class:`.TypeEngine` object
  203. The type object is considered to be significant for cache key
  204. generation.
  205. """
  206. dp_plain_dict = "PD"
  207. """Visit a dictionary with string keys.
  208. The keys of the dictionary should be strings, the values should
  209. be immutable and hashable. The dictionary is considered to be
  210. significant for cache key generation.
  211. """
  212. dp_dialect_options = "DO"
  213. """Visit a dialect options structure."""
  214. dp_string_clauseelement_dict = "CD"
  215. """Visit a dictionary of string keys to :class:`_expression.ClauseElement`
  216. objects.
  217. """
  218. dp_string_multi_dict = "MD"
  219. """Visit a dictionary of string keys to values which may either be
  220. plain immutable/hashable or :class:`.HasCacheKey` objects.
  221. """
  222. dp_annotations_key = "AK"
  223. """Visit the _annotations_cache_key element.
  224. This is a dictionary of additional information about a ClauseElement
  225. that modifies its role. It should be included when comparing or caching
  226. objects, however generating this key is relatively expensive. Visitors
  227. should check the "_annotations" dict for non-None first before creating
  228. this key.
  229. """
  230. dp_plain_obj = "PO"
  231. """Visit a plain python object.
  232. The value should be immutable and hashable, such as an integer.
  233. The value is considered to be significant for cache key generation.
  234. """
  235. dp_named_ddl_element = "DD"
  236. """Visit a simple named DDL element.
  237. The current object used by this method is the :class:`.Sequence`.
  238. The object is only considered to be important for cache key generation
  239. as far as its name, but not any other aspects of it.
  240. """
  241. dp_prefix_sequence = "PS"
  242. """Visit the sequence represented by :class:`_expression.HasPrefixes`
  243. or :class:`_expression.HasSuffixes`.
  244. """
  245. dp_table_hint_list = "TH"
  246. """Visit the ``_hints`` collection of a :class:`_expression.Select`
  247. object.
  248. """
  249. dp_setup_join_tuple = "SJ"
  250. dp_memoized_select_entities = "ME"
  251. dp_statement_hint_list = "SH"
  252. """Visit the ``_statement_hints`` collection of a
  253. :class:`_expression.Select`
  254. object.
  255. """
  256. dp_unknown_structure = "UK"
  257. """Visit an unknown structure.
  258. """
  259. dp_dml_ordered_values = "DML_OV"
  260. """Visit the values() ordered tuple list of an
  261. :class:`_expression.Update` object."""
  262. dp_dml_values = "DML_V"
  263. """Visit the values() dictionary of a :class:`.ValuesBase`
  264. (e.g. Insert or Update) object.
  265. """
  266. dp_dml_multi_values = "DML_MV"
  267. """Visit the values() multi-valued list of dictionaries of an
  268. :class:`_expression.Insert` object.
  269. """
  270. dp_propagate_attrs = "PA"
  271. """Visit the propagate attrs dict. This hardcodes to the particular
  272. elements we care about right now."""
  273. """Symbols that follow are additional symbols that are useful in
  274. caching applications.
  275. Traversals for :class:`_expression.ClauseElement` objects only need to use
  276. those symbols present in :class:`.InternalTraversal`. However, for
  277. additional caching use cases within the ORM, symbols dealing with the
  278. :class:`.HasCacheKey` class are added here.
  279. """
  280. dp_ignore = "IG"
  281. """Specify an object that should be ignored entirely.
  282. This currently applies function call argument caching where some
  283. arguments should not be considered to be part of a cache key.
  284. """
  285. dp_inspectable = "IS"
  286. """Visit an inspectable object where the return value is a
  287. :class:`.HasCacheKey` object."""
  288. dp_multi = "M"
  289. """Visit an object that may be a :class:`.HasCacheKey` or may be a
  290. plain hashable object."""
  291. dp_multi_list = "MT"
  292. """Visit a tuple containing elements that may be :class:`.HasCacheKey` or
  293. may be a plain hashable object."""
  294. dp_has_cache_key_tuples = "HT"
  295. """Visit a list of tuples which contain :class:`.HasCacheKey`
  296. objects.
  297. """
  298. dp_inspectable_list = "IL"
  299. """Visit a list of inspectable objects which upon inspection are
  300. HasCacheKey objects."""
  301. _TraverseInternalsType = List[Tuple[str, InternalTraversal]]
  302. """a structure that defines how a HasTraverseInternals should be
  303. traversed.
  304. This structure consists of a list of (attributename, internaltraversal)
  305. tuples, where the "attributename" refers to the name of an attribute on an
  306. instance of the HasTraverseInternals object, and "internaltraversal" refers
  307. to an :class:`.InternalTraversal` enumeration symbol defining what kind
  308. of data this attribute stores, which indicates to the traverser how it should
  309. be handled.
  310. """
  311. class HasTraverseInternals:
  312. """base for classes that have a "traverse internals" element,
  313. which defines all kinds of ways of traversing the elements of an object.
  314. Compared to :class:`.Visitable`, which relies upon an external visitor to
  315. define how the object is travered (i.e. the :class:`.SQLCompiler`), the
  316. :class:`.HasTraverseInternals` interface allows classes to define their own
  317. traversal, that is, what attributes are accessed and in what order.
  318. """
  319. __slots__ = ()
  320. _traverse_internals: _TraverseInternalsType
  321. _is_immutable: bool = False
  322. @util.preload_module("sqlalchemy.sql.traversals")
  323. def get_children(
  324. self, *, omit_attrs: Tuple[str, ...] = (), **kw: Any
  325. ) -> Iterable[HasTraverseInternals]:
  326. r"""Return immediate child :class:`.visitors.HasTraverseInternals`
  327. elements of this :class:`.visitors.HasTraverseInternals`.
  328. This is used for visit traversal.
  329. \**kw may contain flags that change the collection that is
  330. returned, for example to return a subset of items in order to
  331. cut down on larger traversals, or to return child items from a
  332. different context (such as schema-level collections instead of
  333. clause-level).
  334. """
  335. traversals = util.preloaded.sql_traversals
  336. try:
  337. traverse_internals = self._traverse_internals
  338. except AttributeError:
  339. # user-defined classes may not have a _traverse_internals
  340. return []
  341. dispatch = traversals._get_children.run_generated_dispatch
  342. return itertools.chain.from_iterable(
  343. meth(obj, **kw)
  344. for attrname, obj, meth in dispatch(
  345. self, traverse_internals, "_generated_get_children_traversal"
  346. )
  347. if attrname not in omit_attrs and obj is not None
  348. )
  349. class _InternalTraversalDispatchType(Protocol):
  350. def __call__(s, self: object, visitor: HasTraversalDispatch) -> Any: ...
  351. class HasTraversalDispatch:
  352. r"""Define infrastructure for classes that perform internal traversals
  353. .. versionadded:: 2.0
  354. """
  355. __slots__ = ()
  356. _dispatch_lookup: ClassVar[Dict[Union[InternalTraversal, str], str]] = {}
  357. def dispatch(self, visit_symbol: InternalTraversal) -> Callable[..., Any]:
  358. """Given a method from :class:`.HasTraversalDispatch`, return the
  359. corresponding method on a subclass.
  360. """
  361. name = _dispatch_lookup[visit_symbol]
  362. return getattr(self, name, None) # type: ignore
  363. def run_generated_dispatch(
  364. self,
  365. target: object,
  366. internal_dispatch: _TraverseInternalsType,
  367. generate_dispatcher_name: str,
  368. ) -> Any:
  369. dispatcher: _InternalTraversalDispatchType
  370. try:
  371. dispatcher = target.__class__.__dict__[generate_dispatcher_name]
  372. except KeyError:
  373. # traversals.py -> _preconfigure_traversals()
  374. # may be used to run these ahead of time, but
  375. # is not enabled right now.
  376. # this block will generate any remaining dispatchers.
  377. dispatcher = self.generate_dispatch(
  378. target.__class__, internal_dispatch, generate_dispatcher_name
  379. )
  380. return dispatcher(target, self)
  381. def generate_dispatch(
  382. self,
  383. target_cls: Type[object],
  384. internal_dispatch: _TraverseInternalsType,
  385. generate_dispatcher_name: str,
  386. ) -> _InternalTraversalDispatchType:
  387. dispatcher = self._generate_dispatcher(
  388. internal_dispatch, generate_dispatcher_name
  389. )
  390. # assert isinstance(target_cls, type)
  391. setattr(target_cls, generate_dispatcher_name, dispatcher)
  392. return dispatcher
  393. def _generate_dispatcher(
  394. self, internal_dispatch: _TraverseInternalsType, method_name: str
  395. ) -> _InternalTraversalDispatchType:
  396. names = []
  397. for attrname, visit_sym in internal_dispatch:
  398. meth = self.dispatch(visit_sym)
  399. if meth is not None:
  400. visit_name = _dispatch_lookup[visit_sym]
  401. names.append((attrname, visit_name))
  402. code = (
  403. (" return [\n")
  404. + (
  405. ", \n".join(
  406. " (%r, self.%s, visitor.%s)"
  407. % (attrname, attrname, visit_name)
  408. for attrname, visit_name in names
  409. )
  410. )
  411. + ("\n ]\n")
  412. )
  413. meth_text = ("def %s(self, visitor):\n" % method_name) + code + "\n"
  414. return cast(
  415. _InternalTraversalDispatchType,
  416. langhelpers._exec_code_in_env(meth_text, {}, method_name),
  417. )
  418. ExtendedInternalTraversal = InternalTraversal
  419. def _generate_traversal_dispatch() -> None:
  420. lookup = _dispatch_lookup
  421. for sym in InternalTraversal:
  422. key = sym.name
  423. if key.startswith("dp_"):
  424. visit_key = key.replace("dp_", "visit_")
  425. sym_name = sym.value
  426. assert sym_name not in lookup, sym_name
  427. lookup[sym] = lookup[sym_name] = visit_key
  428. _dispatch_lookup = HasTraversalDispatch._dispatch_lookup
  429. _generate_traversal_dispatch()
  430. class ExternallyTraversible(HasTraverseInternals, Visitable):
  431. __slots__ = ()
  432. _annotations: Mapping[Any, Any] = util.EMPTY_DICT
  433. if typing.TYPE_CHECKING:
  434. def _annotate(self, values: _AnnotationDict) -> Self: ...
  435. def get_children(
  436. self, *, omit_attrs: Tuple[str, ...] = (), **kw: Any
  437. ) -> Iterable[ExternallyTraversible]: ...
  438. def _clone(self, **kw: Any) -> Self:
  439. """clone this element"""
  440. raise NotImplementedError()
  441. def _copy_internals(
  442. self, *, omit_attrs: Tuple[str, ...] = (), **kw: Any
  443. ) -> None:
  444. """Reassign internal elements to be clones of themselves.
  445. Called during a copy-and-traverse operation on newly
  446. shallow-copied elements to create a deep copy.
  447. The given clone function should be used, which may be applying
  448. additional transformations to the element (i.e. replacement
  449. traversal, cloned traversal, annotations).
  450. """
  451. raise NotImplementedError()
  452. _ET = TypeVar("_ET", bound=ExternallyTraversible)
  453. _CE = TypeVar("_CE", bound="ColumnElement[Any]")
  454. _TraverseCallableType = Callable[[_ET], None]
  455. class _CloneCallableType(Protocol):
  456. def __call__(self, element: _ET, **kw: Any) -> _ET: ...
  457. class _TraverseTransformCallableType(Protocol[_ET]):
  458. def __call__(self, element: _ET, **kw: Any) -> Optional[_ET]: ...
  459. _ExtT = TypeVar("_ExtT", bound="ExternalTraversal")
  460. class ExternalTraversal(util.MemoizedSlots):
  461. """Base class for visitor objects which can traverse externally using
  462. the :func:`.visitors.traverse` function.
  463. Direct usage of the :func:`.visitors.traverse` function is usually
  464. preferred.
  465. """
  466. __slots__ = ("_visitor_dict", "_next")
  467. __traverse_options__: Dict[str, Any] = {}
  468. _next: Optional[ExternalTraversal]
  469. def traverse_single(self, obj: Visitable, **kw: Any) -> Any:
  470. for v in self.visitor_iterator:
  471. meth = getattr(v, "visit_%s" % obj.__visit_name__, None)
  472. if meth:
  473. return meth(obj, **kw)
  474. def iterate(
  475. self, obj: Optional[ExternallyTraversible]
  476. ) -> Iterator[ExternallyTraversible]:
  477. """Traverse the given expression structure, returning an iterator
  478. of all elements.
  479. """
  480. return iterate(obj, self.__traverse_options__)
  481. @overload
  482. def traverse(self, obj: Literal[None]) -> None: ...
  483. @overload
  484. def traverse(
  485. self, obj: ExternallyTraversible
  486. ) -> ExternallyTraversible: ...
  487. def traverse(
  488. self, obj: Optional[ExternallyTraversible]
  489. ) -> Optional[ExternallyTraversible]:
  490. """Traverse and visit the given expression structure."""
  491. return traverse(obj, self.__traverse_options__, self._visitor_dict)
  492. def _memoized_attr__visitor_dict(
  493. self,
  494. ) -> Dict[str, _TraverseCallableType[Any]]:
  495. visitors = {}
  496. for name in dir(self):
  497. if name.startswith("visit_"):
  498. visitors[name[6:]] = getattr(self, name)
  499. return visitors
  500. @property
  501. def visitor_iterator(self) -> Iterator[ExternalTraversal]:
  502. """Iterate through this visitor and each 'chained' visitor."""
  503. v: Optional[ExternalTraversal] = self
  504. while v:
  505. yield v
  506. v = getattr(v, "_next", None)
  507. def chain(self: _ExtT, visitor: ExternalTraversal) -> _ExtT:
  508. """'Chain' an additional ExternalTraversal onto this ExternalTraversal
  509. The chained visitor will receive all visit events after this one.
  510. """
  511. tail = list(self.visitor_iterator)[-1]
  512. tail._next = visitor
  513. return self
  514. class CloningExternalTraversal(ExternalTraversal):
  515. """Base class for visitor objects which can traverse using
  516. the :func:`.visitors.cloned_traverse` function.
  517. Direct usage of the :func:`.visitors.cloned_traverse` function is usually
  518. preferred.
  519. """
  520. __slots__ = ()
  521. def copy_and_process(
  522. self, list_: List[ExternallyTraversible]
  523. ) -> List[ExternallyTraversible]:
  524. """Apply cloned traversal to the given list of elements, and return
  525. the new list.
  526. """
  527. return [self.traverse(x) for x in list_]
  528. @overload
  529. def traverse(self, obj: Literal[None]) -> None: ...
  530. @overload
  531. def traverse(
  532. self, obj: ExternallyTraversible
  533. ) -> ExternallyTraversible: ...
  534. def traverse(
  535. self, obj: Optional[ExternallyTraversible]
  536. ) -> Optional[ExternallyTraversible]:
  537. """Traverse and visit the given expression structure."""
  538. return cloned_traverse(
  539. obj, self.__traverse_options__, self._visitor_dict
  540. )
  541. class ReplacingExternalTraversal(CloningExternalTraversal):
  542. """Base class for visitor objects which can traverse using
  543. the :func:`.visitors.replacement_traverse` function.
  544. Direct usage of the :func:`.visitors.replacement_traverse` function is
  545. usually preferred.
  546. """
  547. __slots__ = ()
  548. def replace(
  549. self, elem: ExternallyTraversible
  550. ) -> Optional[ExternallyTraversible]:
  551. """Receive pre-copied elements during a cloning traversal.
  552. If the method returns a new element, the element is used
  553. instead of creating a simple copy of the element. Traversal
  554. will halt on the newly returned element if it is re-encountered.
  555. """
  556. return None
  557. @overload
  558. def traverse(self, obj: Literal[None]) -> None: ...
  559. @overload
  560. def traverse(
  561. self, obj: ExternallyTraversible
  562. ) -> ExternallyTraversible: ...
  563. def traverse(
  564. self, obj: Optional[ExternallyTraversible]
  565. ) -> Optional[ExternallyTraversible]:
  566. """Traverse and visit the given expression structure."""
  567. def replace(
  568. element: ExternallyTraversible,
  569. **kw: Any,
  570. ) -> Optional[ExternallyTraversible]:
  571. for v in self.visitor_iterator:
  572. e = cast(ReplacingExternalTraversal, v).replace(element)
  573. if e is not None:
  574. return e
  575. return None
  576. return replacement_traverse(obj, self.__traverse_options__, replace)
  577. # backwards compatibility
  578. Traversible = Visitable
  579. ClauseVisitor = ExternalTraversal
  580. CloningVisitor = CloningExternalTraversal
  581. ReplacingCloningVisitor = ReplacingExternalTraversal
  582. def iterate(
  583. obj: Optional[ExternallyTraversible],
  584. opts: Mapping[str, Any] = util.EMPTY_DICT,
  585. ) -> Iterator[ExternallyTraversible]:
  586. r"""Traverse the given expression structure, returning an iterator.
  587. Traversal is configured to be breadth-first.
  588. The central API feature used by the :func:`.visitors.iterate`
  589. function is the
  590. :meth:`_expression.ClauseElement.get_children` method of
  591. :class:`_expression.ClauseElement` objects. This method should return all
  592. the :class:`_expression.ClauseElement` objects which are associated with a
  593. particular :class:`_expression.ClauseElement` object. For example, a
  594. :class:`.Case` structure will refer to a series of
  595. :class:`_expression.ColumnElement` objects within its "whens" and "else\_"
  596. member variables.
  597. :param obj: :class:`_expression.ClauseElement` structure to be traversed
  598. :param opts: dictionary of iteration options. This dictionary is usually
  599. empty in modern usage.
  600. """
  601. if obj is None:
  602. return
  603. yield obj
  604. children = obj.get_children(**opts)
  605. if not children:
  606. return
  607. stack = deque([children])
  608. while stack:
  609. t_iterator = stack.popleft()
  610. for t in t_iterator:
  611. yield t
  612. stack.append(t.get_children(**opts))
  613. @overload
  614. def traverse_using(
  615. iterator: Iterable[ExternallyTraversible],
  616. obj: Literal[None],
  617. visitors: Mapping[str, _TraverseCallableType[Any]],
  618. ) -> None: ...
  619. @overload
  620. def traverse_using(
  621. iterator: Iterable[ExternallyTraversible],
  622. obj: ExternallyTraversible,
  623. visitors: Mapping[str, _TraverseCallableType[Any]],
  624. ) -> ExternallyTraversible: ...
  625. def traverse_using(
  626. iterator: Iterable[ExternallyTraversible],
  627. obj: Optional[ExternallyTraversible],
  628. visitors: Mapping[str, _TraverseCallableType[Any]],
  629. ) -> Optional[ExternallyTraversible]:
  630. """Visit the given expression structure using the given iterator of
  631. objects.
  632. :func:`.visitors.traverse_using` is usually called internally as the result
  633. of the :func:`.visitors.traverse` function.
  634. :param iterator: an iterable or sequence which will yield
  635. :class:`_expression.ClauseElement`
  636. structures; the iterator is assumed to be the
  637. product of the :func:`.visitors.iterate` function.
  638. :param obj: the :class:`_expression.ClauseElement`
  639. that was used as the target of the
  640. :func:`.iterate` function.
  641. :param visitors: dictionary of visit functions. See :func:`.traverse`
  642. for details on this dictionary.
  643. .. seealso::
  644. :func:`.traverse`
  645. """
  646. for target in iterator:
  647. meth = visitors.get(target.__visit_name__, None)
  648. if meth:
  649. meth(target)
  650. return obj
  651. @overload
  652. def traverse(
  653. obj: Literal[None],
  654. opts: Mapping[str, Any],
  655. visitors: Mapping[str, _TraverseCallableType[Any]],
  656. ) -> None: ...
  657. @overload
  658. def traverse(
  659. obj: ExternallyTraversible,
  660. opts: Mapping[str, Any],
  661. visitors: Mapping[str, _TraverseCallableType[Any]],
  662. ) -> ExternallyTraversible: ...
  663. def traverse(
  664. obj: Optional[ExternallyTraversible],
  665. opts: Mapping[str, Any],
  666. visitors: Mapping[str, _TraverseCallableType[Any]],
  667. ) -> Optional[ExternallyTraversible]:
  668. """Traverse and visit the given expression structure using the default
  669. iterator.
  670. e.g.::
  671. from sqlalchemy.sql import visitors
  672. stmt = select(some_table).where(some_table.c.foo == "bar")
  673. def visit_bindparam(bind_param):
  674. print("found bound value: %s" % bind_param.value)
  675. visitors.traverse(stmt, {}, {"bindparam": visit_bindparam})
  676. The iteration of objects uses the :func:`.visitors.iterate` function,
  677. which does a breadth-first traversal using a stack.
  678. :param obj: :class:`_expression.ClauseElement` structure to be traversed
  679. :param opts: dictionary of iteration options. This dictionary is usually
  680. empty in modern usage.
  681. :param visitors: dictionary of visit functions. The dictionary should
  682. have strings as keys, each of which would correspond to the
  683. ``__visit_name__`` of a particular kind of SQL expression object, and
  684. callable functions as values, each of which represents a visitor function
  685. for that kind of object.
  686. """
  687. return traverse_using(iterate(obj, opts), obj, visitors)
  688. @overload
  689. def cloned_traverse(
  690. obj: Literal[None],
  691. opts: Mapping[str, Any],
  692. visitors: Mapping[str, _TraverseCallableType[Any]],
  693. ) -> None: ...
  694. # a bit of controversy here, as the clone of the lead element
  695. # *could* in theory replace with an entirely different kind of element.
  696. # however this is really not how cloned_traverse is ever used internally
  697. # at least.
  698. @overload
  699. def cloned_traverse(
  700. obj: _ET,
  701. opts: Mapping[str, Any],
  702. visitors: Mapping[str, _TraverseCallableType[Any]],
  703. ) -> _ET: ...
  704. def cloned_traverse(
  705. obj: Optional[ExternallyTraversible],
  706. opts: Mapping[str, Any],
  707. visitors: Mapping[str, _TraverseCallableType[Any]],
  708. ) -> Optional[ExternallyTraversible]:
  709. """Clone the given expression structure, allowing modifications by
  710. visitors for mutable objects.
  711. Traversal usage is the same as that of :func:`.visitors.traverse`.
  712. The visitor functions present in the ``visitors`` dictionary may also
  713. modify the internals of the given structure as the traversal proceeds.
  714. The :func:`.cloned_traverse` function does **not** provide objects that are
  715. part of the :class:`.Immutable` interface to the visit methods (this
  716. primarily includes :class:`.ColumnClause`, :class:`.Column`,
  717. :class:`.TableClause` and :class:`.Table` objects). As this traversal is
  718. only intended to allow in-place mutation of objects, :class:`.Immutable`
  719. objects are skipped. The :meth:`.Immutable._clone` method is still called
  720. on each object to allow for objects to replace themselves with a different
  721. object based on a clone of their sub-internals (e.g. a
  722. :class:`.ColumnClause` that clones its subquery to return a new
  723. :class:`.ColumnClause`).
  724. .. versionchanged:: 2.0 The :func:`.cloned_traverse` function omits
  725. objects that are part of the :class:`.Immutable` interface.
  726. The central API feature used by the :func:`.visitors.cloned_traverse`
  727. and :func:`.visitors.replacement_traverse` functions, in addition to the
  728. :meth:`_expression.ClauseElement.get_children`
  729. function that is used to achieve
  730. the iteration, is the :meth:`_expression.ClauseElement._copy_internals`
  731. method.
  732. For a :class:`_expression.ClauseElement`
  733. structure to support cloning and replacement
  734. traversals correctly, it needs to be able to pass a cloning function into
  735. its internal members in order to make copies of them.
  736. .. seealso::
  737. :func:`.visitors.traverse`
  738. :func:`.visitors.replacement_traverse`
  739. """
  740. cloned: Dict[int, ExternallyTraversible] = {}
  741. stop_on = set(opts.get("stop_on", []))
  742. def deferred_copy_internals(
  743. obj: ExternallyTraversible,
  744. ) -> ExternallyTraversible:
  745. return cloned_traverse(obj, opts, visitors)
  746. def clone(elem: ExternallyTraversible, **kw: Any) -> ExternallyTraversible:
  747. if elem in stop_on:
  748. return elem
  749. else:
  750. if id(elem) not in cloned:
  751. if "replace" in kw:
  752. newelem = cast(
  753. Optional[ExternallyTraversible], kw["replace"](elem)
  754. )
  755. if newelem is not None:
  756. cloned[id(elem)] = newelem
  757. return newelem
  758. # the _clone method for immutable normally returns "self".
  759. # however, the method is still allowed to return a
  760. # different object altogether; ColumnClause._clone() will
  761. # based on options clone the subquery to which it is associated
  762. # and return the new corresponding column.
  763. cloned[id(elem)] = newelem = elem._clone(clone=clone, **kw)
  764. newelem._copy_internals(clone=clone, **kw)
  765. # however, visit methods which are tasked with in-place
  766. # mutation of the object should not get access to the immutable
  767. # object.
  768. if not elem._is_immutable:
  769. meth = visitors.get(newelem.__visit_name__, None)
  770. if meth:
  771. meth(newelem)
  772. return cloned[id(elem)]
  773. if obj is not None:
  774. obj = clone(
  775. obj, deferred_copy_internals=deferred_copy_internals, **opts
  776. )
  777. clone = None # type: ignore[assignment] # remove gc cycles
  778. return obj
  779. @overload
  780. def replacement_traverse(
  781. obj: Literal[None],
  782. opts: Mapping[str, Any],
  783. replace: _TraverseTransformCallableType[Any],
  784. ) -> None: ...
  785. @overload
  786. def replacement_traverse(
  787. obj: _CE,
  788. opts: Mapping[str, Any],
  789. replace: _TraverseTransformCallableType[Any],
  790. ) -> _CE: ...
  791. @overload
  792. def replacement_traverse(
  793. obj: ExternallyTraversible,
  794. opts: Mapping[str, Any],
  795. replace: _TraverseTransformCallableType[Any],
  796. ) -> ExternallyTraversible: ...
  797. def replacement_traverse(
  798. obj: Optional[ExternallyTraversible],
  799. opts: Mapping[str, Any],
  800. replace: _TraverseTransformCallableType[Any],
  801. ) -> Optional[ExternallyTraversible]:
  802. """Clone the given expression structure, allowing element
  803. replacement by a given replacement function.
  804. This function is very similar to the :func:`.visitors.cloned_traverse`
  805. function, except instead of being passed a dictionary of visitors, all
  806. elements are unconditionally passed into the given replace function.
  807. The replace function then has the option to return an entirely new object
  808. which will replace the one given. If it returns ``None``, then the object
  809. is kept in place.
  810. The difference in usage between :func:`.visitors.cloned_traverse` and
  811. :func:`.visitors.replacement_traverse` is that in the former case, an
  812. already-cloned object is passed to the visitor function, and the visitor
  813. function can then manipulate the internal state of the object.
  814. In the case of the latter, the visitor function should only return an
  815. entirely different object, or do nothing.
  816. The use case for :func:`.visitors.replacement_traverse` is that of
  817. replacing a FROM clause inside of a SQL structure with a different one,
  818. as is a common use case within the ORM.
  819. """
  820. cloned = {}
  821. stop_on = {id(x) for x in opts.get("stop_on", [])}
  822. def deferred_copy_internals(
  823. obj: ExternallyTraversible,
  824. ) -> ExternallyTraversible:
  825. return replacement_traverse(obj, opts, replace)
  826. def clone(elem: ExternallyTraversible, **kw: Any) -> ExternallyTraversible:
  827. if (
  828. id(elem) in stop_on
  829. or "no_replacement_traverse" in elem._annotations
  830. ):
  831. return elem
  832. else:
  833. newelem = replace(elem)
  834. if newelem is not None:
  835. stop_on.add(id(newelem))
  836. return newelem # type: ignore
  837. else:
  838. # base "already seen" on id(), not hash, so that we don't
  839. # replace an Annotated element with its non-annotated one, and
  840. # vice versa
  841. id_elem = id(elem)
  842. if id_elem not in cloned:
  843. if "replace" in kw:
  844. newelem = kw["replace"](elem)
  845. if newelem is not None:
  846. cloned[id_elem] = newelem
  847. return newelem # type: ignore
  848. cloned[id_elem] = newelem = elem._clone(**kw)
  849. newelem._copy_internals(clone=clone, **kw)
  850. return cloned[id_elem] # type: ignore
  851. if obj is not None:
  852. obj = clone(
  853. obj, deferred_copy_internals=deferred_copy_internals, **opts
  854. )
  855. clone = None # type: ignore[assignment] # remove gc cycles
  856. return obj