pipeline.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663
  1. """Experimental pipeline API functionality. Be careful with this API, it's subject to change."""
  2. from __future__ import annotations
  3. import datetime
  4. import operator
  5. import re
  6. import sys
  7. from collections import deque
  8. from collections.abc import Container
  9. from dataclasses import dataclass
  10. from decimal import Decimal
  11. from functools import cached_property, partial
  12. from re import Pattern
  13. from typing import TYPE_CHECKING, Annotated, Any, Callable, Generic, Protocol, TypeVar, Union, overload
  14. import annotated_types
  15. if TYPE_CHECKING:
  16. from pydantic import GetCoreSchemaHandler
  17. from pydantic_core import PydanticCustomError
  18. from pydantic_core import core_schema as cs
  19. from pydantic import Strict
  20. from pydantic._internal._internal_dataclass import slots_true as _slots_true
  21. if sys.version_info < (3, 10):
  22. EllipsisType = type(Ellipsis)
  23. else:
  24. from types import EllipsisType
  25. __all__ = ['validate_as', 'validate_as_deferred', 'transform']
  26. _slots_frozen = {**_slots_true, 'frozen': True}
  27. @dataclass(**_slots_frozen)
  28. class _ValidateAs:
  29. tp: type[Any]
  30. strict: bool = False
  31. @dataclass
  32. class _ValidateAsDefer:
  33. func: Callable[[], type[Any]]
  34. @cached_property
  35. def tp(self) -> type[Any]:
  36. return self.func()
  37. @dataclass(**_slots_frozen)
  38. class _Transform:
  39. func: Callable[[Any], Any]
  40. @dataclass(**_slots_frozen)
  41. class _PipelineOr:
  42. left: _Pipeline[Any, Any]
  43. right: _Pipeline[Any, Any]
  44. @dataclass(**_slots_frozen)
  45. class _PipelineAnd:
  46. left: _Pipeline[Any, Any]
  47. right: _Pipeline[Any, Any]
  48. @dataclass(**_slots_frozen)
  49. class _Eq:
  50. value: Any
  51. @dataclass(**_slots_frozen)
  52. class _NotEq:
  53. value: Any
  54. @dataclass(**_slots_frozen)
  55. class _In:
  56. values: Container[Any]
  57. @dataclass(**_slots_frozen)
  58. class _NotIn:
  59. values: Container[Any]
  60. _ConstraintAnnotation = Union[
  61. annotated_types.Le,
  62. annotated_types.Ge,
  63. annotated_types.Lt,
  64. annotated_types.Gt,
  65. annotated_types.Len,
  66. annotated_types.MultipleOf,
  67. annotated_types.Timezone,
  68. annotated_types.Interval,
  69. annotated_types.Predicate,
  70. # common predicates not included in annotated_types
  71. _Eq,
  72. _NotEq,
  73. _In,
  74. _NotIn,
  75. # regular expressions
  76. Pattern[str],
  77. ]
  78. @dataclass(**_slots_frozen)
  79. class _Constraint:
  80. constraint: _ConstraintAnnotation
  81. _Step = Union[_ValidateAs, _ValidateAsDefer, _Transform, _PipelineOr, _PipelineAnd, _Constraint]
  82. _InT = TypeVar('_InT')
  83. _OutT = TypeVar('_OutT')
  84. _NewOutT = TypeVar('_NewOutT')
  85. class _FieldTypeMarker:
  86. pass
  87. # TODO: ultimately, make this public, see https://github.com/pydantic/pydantic/pull/9459#discussion_r1628197626
  88. # Also, make this frozen eventually, but that doesn't work right now because of the generic base
  89. # Which attempts to modify __orig_base__ and such.
  90. # We could go with a manual freeze, but that seems overkill for now.
  91. @dataclass(**_slots_true)
  92. class _Pipeline(Generic[_InT, _OutT]):
  93. """Abstract representation of a chain of validation, transformation, and parsing steps."""
  94. _steps: tuple[_Step, ...]
  95. def transform(
  96. self,
  97. func: Callable[[_OutT], _NewOutT],
  98. ) -> _Pipeline[_InT, _NewOutT]:
  99. """Transform the output of the previous step.
  100. If used as the first step in a pipeline, the type of the field is used.
  101. That is, the transformation is applied to after the value is parsed to the field's type.
  102. """
  103. return _Pipeline[_InT, _NewOutT](self._steps + (_Transform(func),))
  104. @overload
  105. def validate_as(self, tp: type[_NewOutT], *, strict: bool = False) -> _Pipeline[_InT, _NewOutT]: ...
  106. @overload
  107. def validate_as(
  108. self,
  109. tp: ellipsis, # noqa: F821 # TODO: use `_typing_extra.EllipsisType` when we drop Py3.9
  110. *,
  111. strict: bool = False,
  112. ) -> _Pipeline[_InT, Any]: ...
  113. # TODO PEP 747: use TypeForm to properly type Annotated aliases (e.g. NewPath, FilePath).
  114. # This fallback accepts any type expression but loses generic type inference.
  115. @overload
  116. def validate_as(self, tp: Any, *, strict: bool = ...) -> _Pipeline[_InT, Any]: ...
  117. def validate_as(self, tp: type[_NewOutT] | EllipsisType | Any, *, strict: bool = False) -> _Pipeline[_InT, Any]: # type: ignore
  118. """Validate / parse the input into a new type.
  119. If no type is provided, the type of the field is used.
  120. Types are parsed in Pydantic's `lax` mode by default,
  121. but you can enable `strict` mode by passing `strict=True`.
  122. """
  123. if isinstance(tp, EllipsisType):
  124. return _Pipeline[_InT, Any](self._steps + (_ValidateAs(_FieldTypeMarker, strict=strict),))
  125. return _Pipeline[_InT, _NewOutT](self._steps + (_ValidateAs(tp, strict=strict),))
  126. def validate_as_deferred(self, func: Callable[[], type[_NewOutT]]) -> _Pipeline[_InT, _NewOutT]:
  127. """Parse the input into a new type, deferring resolution of the type until the current class
  128. is fully defined.
  129. This is useful when you need to reference the class in it's own type annotations.
  130. """
  131. return _Pipeline[_InT, _NewOutT](self._steps + (_ValidateAsDefer(func),))
  132. # constraints
  133. @overload
  134. def constrain(self: _Pipeline[_InT, _NewOutGe], constraint: annotated_types.Ge) -> _Pipeline[_InT, _NewOutGe]: ...
  135. @overload
  136. def constrain(self: _Pipeline[_InT, _NewOutGt], constraint: annotated_types.Gt) -> _Pipeline[_InT, _NewOutGt]: ...
  137. @overload
  138. def constrain(self: _Pipeline[_InT, _NewOutLe], constraint: annotated_types.Le) -> _Pipeline[_InT, _NewOutLe]: ...
  139. @overload
  140. def constrain(self: _Pipeline[_InT, _NewOutLt], constraint: annotated_types.Lt) -> _Pipeline[_InT, _NewOutLt]: ...
  141. @overload
  142. def constrain(
  143. self: _Pipeline[_InT, _NewOutLen], constraint: annotated_types.Len
  144. ) -> _Pipeline[_InT, _NewOutLen]: ...
  145. @overload
  146. def constrain(
  147. self: _Pipeline[_InT, _NewOutT], constraint: annotated_types.MultipleOf
  148. ) -> _Pipeline[_InT, _NewOutT]: ...
  149. @overload
  150. def constrain(
  151. self: _Pipeline[_InT, _NewOutDatetime], constraint: annotated_types.Timezone
  152. ) -> _Pipeline[_InT, _NewOutDatetime]: ...
  153. @overload
  154. def constrain(self: _Pipeline[_InT, _OutT], constraint: annotated_types.Predicate) -> _Pipeline[_InT, _OutT]: ...
  155. @overload
  156. def constrain(
  157. self: _Pipeline[_InT, _NewOutInterval], constraint: annotated_types.Interval
  158. ) -> _Pipeline[_InT, _NewOutInterval]: ...
  159. @overload
  160. def constrain(self: _Pipeline[_InT, _OutT], constraint: _Eq) -> _Pipeline[_InT, _OutT]: ...
  161. @overload
  162. def constrain(self: _Pipeline[_InT, _OutT], constraint: _NotEq) -> _Pipeline[_InT, _OutT]: ...
  163. @overload
  164. def constrain(self: _Pipeline[_InT, _OutT], constraint: _In) -> _Pipeline[_InT, _OutT]: ...
  165. @overload
  166. def constrain(self: _Pipeline[_InT, _OutT], constraint: _NotIn) -> _Pipeline[_InT, _OutT]: ...
  167. @overload
  168. def constrain(self: _Pipeline[_InT, _NewOutT], constraint: Pattern[str]) -> _Pipeline[_InT, _NewOutT]: ...
  169. def constrain(self, constraint: _ConstraintAnnotation) -> Any:
  170. """Constrain a value to meet a certain condition.
  171. We support most conditions from `annotated_types`, as well as regular expressions.
  172. Most of the time you'll be calling a shortcut method like `gt`, `lt`, `len`, etc
  173. so you don't need to call this directly.
  174. """
  175. return _Pipeline[_InT, _OutT](self._steps + (_Constraint(constraint),))
  176. def predicate(self: _Pipeline[_InT, _NewOutT], func: Callable[[_NewOutT], bool]) -> _Pipeline[_InT, _NewOutT]:
  177. """Constrain a value to meet a certain predicate."""
  178. return self.constrain(annotated_types.Predicate(func))
  179. def gt(self: _Pipeline[_InT, _NewOutGt], gt: _NewOutGt) -> _Pipeline[_InT, _NewOutGt]:
  180. """Constrain a value to be greater than a certain value."""
  181. return self.constrain(annotated_types.Gt(gt))
  182. def lt(self: _Pipeline[_InT, _NewOutLt], lt: _NewOutLt) -> _Pipeline[_InT, _NewOutLt]:
  183. """Constrain a value to be less than a certain value."""
  184. return self.constrain(annotated_types.Lt(lt))
  185. def ge(self: _Pipeline[_InT, _NewOutGe], ge: _NewOutGe) -> _Pipeline[_InT, _NewOutGe]:
  186. """Constrain a value to be greater than or equal to a certain value."""
  187. return self.constrain(annotated_types.Ge(ge))
  188. def le(self: _Pipeline[_InT, _NewOutLe], le: _NewOutLe) -> _Pipeline[_InT, _NewOutLe]:
  189. """Constrain a value to be less than or equal to a certain value."""
  190. return self.constrain(annotated_types.Le(le))
  191. def len(self: _Pipeline[_InT, _NewOutLen], min_len: int, max_len: int | None = None) -> _Pipeline[_InT, _NewOutLen]:
  192. """Constrain a value to have a certain length."""
  193. return self.constrain(annotated_types.Len(min_len, max_len))
  194. @overload
  195. def multiple_of(self: _Pipeline[_InT, _NewOutDiv], multiple_of: _NewOutDiv) -> _Pipeline[_InT, _NewOutDiv]: ...
  196. @overload
  197. def multiple_of(self: _Pipeline[_InT, _NewOutMod], multiple_of: _NewOutMod) -> _Pipeline[_InT, _NewOutMod]: ...
  198. def multiple_of(self: _Pipeline[_InT, Any], multiple_of: Any) -> _Pipeline[_InT, Any]:
  199. """Constrain a value to be a multiple of a certain number."""
  200. return self.constrain(annotated_types.MultipleOf(multiple_of))
  201. def eq(self: _Pipeline[_InT, _OutT], value: _OutT) -> _Pipeline[_InT, _OutT]:
  202. """Constrain a value to be equal to a certain value."""
  203. return self.constrain(_Eq(value))
  204. def not_eq(self: _Pipeline[_InT, _OutT], value: _OutT) -> _Pipeline[_InT, _OutT]:
  205. """Constrain a value to not be equal to a certain value."""
  206. return self.constrain(_NotEq(value))
  207. def in_(self: _Pipeline[_InT, _OutT], values: Container[_OutT]) -> _Pipeline[_InT, _OutT]:
  208. """Constrain a value to be in a certain set."""
  209. return self.constrain(_In(values))
  210. def not_in(self: _Pipeline[_InT, _OutT], values: Container[_OutT]) -> _Pipeline[_InT, _OutT]:
  211. """Constrain a value to not be in a certain set."""
  212. return self.constrain(_NotIn(values))
  213. # timezone methods
  214. def datetime_tz_naive(self: _Pipeline[_InT, datetime.datetime]) -> _Pipeline[_InT, datetime.datetime]:
  215. return self.constrain(annotated_types.Timezone(None))
  216. def datetime_tz_aware(self: _Pipeline[_InT, datetime.datetime]) -> _Pipeline[_InT, datetime.datetime]:
  217. return self.constrain(annotated_types.Timezone(...))
  218. def datetime_tz(
  219. self: _Pipeline[_InT, datetime.datetime], tz: datetime.tzinfo
  220. ) -> _Pipeline[_InT, datetime.datetime]:
  221. return self.constrain(annotated_types.Timezone(tz)) # type: ignore
  222. def datetime_with_tz(
  223. self: _Pipeline[_InT, datetime.datetime], tz: datetime.tzinfo | None
  224. ) -> _Pipeline[_InT, datetime.datetime]:
  225. return self.transform(partial(datetime.datetime.replace, tzinfo=tz))
  226. # string methods
  227. def str_lower(self: _Pipeline[_InT, str]) -> _Pipeline[_InT, str]:
  228. return self.transform(str.lower)
  229. def str_upper(self: _Pipeline[_InT, str]) -> _Pipeline[_InT, str]:
  230. return self.transform(str.upper)
  231. def str_title(self: _Pipeline[_InT, str]) -> _Pipeline[_InT, str]:
  232. return self.transform(str.title)
  233. def str_strip(self: _Pipeline[_InT, str]) -> _Pipeline[_InT, str]:
  234. return self.transform(str.strip)
  235. def str_pattern(self: _Pipeline[_InT, str], pattern: str) -> _Pipeline[_InT, str]:
  236. return self.constrain(re.compile(pattern))
  237. def str_contains(self: _Pipeline[_InT, str], substring: str) -> _Pipeline[_InT, str]:
  238. return self.predicate(lambda v: substring in v)
  239. def str_starts_with(self: _Pipeline[_InT, str], prefix: str) -> _Pipeline[_InT, str]:
  240. return self.predicate(lambda v: v.startswith(prefix))
  241. def str_ends_with(self: _Pipeline[_InT, str], suffix: str) -> _Pipeline[_InT, str]:
  242. return self.predicate(lambda v: v.endswith(suffix))
  243. # operators
  244. def otherwise(self, other: _Pipeline[_OtherIn, _OtherOut]) -> _Pipeline[_InT | _OtherIn, _OutT | _OtherOut]:
  245. """Combine two validation chains, returning the result of the first chain if it succeeds, and the second chain if it fails."""
  246. return _Pipeline((_PipelineOr(self, other),))
  247. __or__ = otherwise
  248. def then(self, other: _Pipeline[_OutT, _OtherOut]) -> _Pipeline[_InT, _OtherOut]:
  249. """Pipe the result of one validation chain into another."""
  250. return _Pipeline((_PipelineAnd(self, other),))
  251. __and__ = then
  252. def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaHandler) -> cs.CoreSchema:
  253. queue = deque(self._steps)
  254. s = None
  255. while queue:
  256. step = queue.popleft()
  257. s = _apply_step(step, s, handler, source_type)
  258. s = s or cs.any_schema()
  259. return s
  260. def __supports_type__(self, _: _OutT) -> bool:
  261. raise NotImplementedError
  262. validate_as = _Pipeline[Any, Any](()).validate_as
  263. validate_as_deferred = _Pipeline[Any, Any](()).validate_as_deferred
  264. transform = _Pipeline[Any, Any]((_ValidateAs(_FieldTypeMarker),)).transform
  265. def _check_func(
  266. func: Callable[[Any], bool], predicate_err: str | Callable[[], str], s: cs.CoreSchema | None
  267. ) -> cs.CoreSchema:
  268. def handler(v: Any) -> Any:
  269. if func(v):
  270. return v
  271. raise ValueError(f'Expected {predicate_err if isinstance(predicate_err, str) else predicate_err()}')
  272. if s is None:
  273. return cs.no_info_plain_validator_function(handler)
  274. else:
  275. return cs.no_info_after_validator_function(handler, s)
  276. def _apply_step(step: _Step, s: cs.CoreSchema | None, handler: GetCoreSchemaHandler, source_type: Any) -> cs.CoreSchema:
  277. if isinstance(step, _ValidateAs):
  278. s = _apply_parse(s, step.tp, step.strict, handler, source_type)
  279. elif isinstance(step, _ValidateAsDefer):
  280. s = _apply_parse(s, step.tp, False, handler, source_type)
  281. elif isinstance(step, _Transform):
  282. s = _apply_transform(s, step.func, handler)
  283. elif isinstance(step, _Constraint):
  284. s = _apply_constraint(s, step.constraint)
  285. elif isinstance(step, _PipelineOr):
  286. s = cs.union_schema([handler(step.left), handler(step.right)])
  287. else:
  288. assert isinstance(step, _PipelineAnd)
  289. s = cs.chain_schema([handler(step.left), handler(step.right)])
  290. return s
  291. def _apply_parse(
  292. s: cs.CoreSchema | None,
  293. tp: type[Any],
  294. strict: bool,
  295. handler: GetCoreSchemaHandler,
  296. source_type: Any,
  297. ) -> cs.CoreSchema:
  298. if tp is _FieldTypeMarker:
  299. return cs.chain_schema([s, handler(source_type)]) if s else handler(source_type)
  300. if strict:
  301. tp = Annotated[tp, Strict()] # type: ignore
  302. if s and s['type'] == 'any':
  303. return handler(tp)
  304. else:
  305. return cs.chain_schema([s, handler(tp)]) if s else handler(tp)
  306. def _apply_transform(
  307. s: cs.CoreSchema | None, func: Callable[[Any], Any], handler: GetCoreSchemaHandler
  308. ) -> cs.CoreSchema:
  309. if s is None:
  310. return cs.no_info_plain_validator_function(func)
  311. if s['type'] == 'str':
  312. if func is str.strip:
  313. s = s.copy()
  314. s['strip_whitespace'] = True
  315. return s
  316. elif func is str.lower:
  317. s = s.copy()
  318. s['to_lower'] = True
  319. return s
  320. elif func is str.upper:
  321. s = s.copy()
  322. s['to_upper'] = True
  323. return s
  324. return cs.no_info_after_validator_function(func, s)
  325. def _apply_constraint( # noqa: C901
  326. s: cs.CoreSchema | None, constraint: _ConstraintAnnotation
  327. ) -> cs.CoreSchema:
  328. """Apply a single constraint to a schema."""
  329. if isinstance(constraint, annotated_types.Gt):
  330. gt = constraint.gt
  331. if s and s['type'] in {'int', 'float', 'decimal'}:
  332. s = s.copy()
  333. if s['type'] == 'int' and isinstance(gt, int):
  334. s['gt'] = gt
  335. elif s['type'] == 'float' and isinstance(gt, float):
  336. s['gt'] = gt
  337. elif s['type'] == 'decimal' and isinstance(gt, Decimal):
  338. s['gt'] = gt
  339. else:
  340. def check_gt(v: Any) -> bool:
  341. return v > gt
  342. s = _check_func(check_gt, f'> {gt}', s)
  343. elif isinstance(constraint, annotated_types.Ge):
  344. ge = constraint.ge
  345. if s and s['type'] in {'int', 'float', 'decimal'}:
  346. s = s.copy()
  347. if s['type'] == 'int' and isinstance(ge, int):
  348. s['ge'] = ge
  349. elif s['type'] == 'float' and isinstance(ge, float):
  350. s['ge'] = ge
  351. elif s['type'] == 'decimal' and isinstance(ge, Decimal):
  352. s['ge'] = ge
  353. def check_ge(v: Any) -> bool:
  354. return v >= ge
  355. s = _check_func(check_ge, f'>= {ge}', s)
  356. elif isinstance(constraint, annotated_types.Lt):
  357. lt = constraint.lt
  358. if s and s['type'] in {'int', 'float', 'decimal'}:
  359. s = s.copy()
  360. if s['type'] == 'int' and isinstance(lt, int):
  361. s['lt'] = lt
  362. elif s['type'] == 'float' and isinstance(lt, float):
  363. s['lt'] = lt
  364. elif s['type'] == 'decimal' and isinstance(lt, Decimal):
  365. s['lt'] = lt
  366. def check_lt(v: Any) -> bool:
  367. return v < lt
  368. s = _check_func(check_lt, f'< {lt}', s)
  369. elif isinstance(constraint, annotated_types.Le):
  370. le = constraint.le
  371. if s and s['type'] in {'int', 'float', 'decimal'}:
  372. s = s.copy()
  373. if s['type'] == 'int' and isinstance(le, int):
  374. s['le'] = le
  375. elif s['type'] == 'float' and isinstance(le, float):
  376. s['le'] = le
  377. elif s['type'] == 'decimal' and isinstance(le, Decimal):
  378. s['le'] = le
  379. def check_le(v: Any) -> bool:
  380. return v <= le
  381. s = _check_func(check_le, f'<= {le}', s)
  382. elif isinstance(constraint, annotated_types.Len):
  383. min_len = constraint.min_length
  384. max_len = constraint.max_length
  385. if s and s['type'] in {'str', 'list', 'tuple', 'set', 'frozenset', 'dict'}:
  386. assert (
  387. s['type'] == 'str'
  388. or s['type'] == 'list'
  389. or s['type'] == 'tuple'
  390. or s['type'] == 'set'
  391. or s['type'] == 'dict'
  392. or s['type'] == 'frozenset'
  393. )
  394. s = s.copy()
  395. if min_len != 0:
  396. s['min_length'] = min_len
  397. if max_len is not None:
  398. s['max_length'] = max_len
  399. def check_len(v: Any) -> bool:
  400. if max_len is not None:
  401. return (min_len <= len(v)) and (len(v) <= max_len)
  402. return min_len <= len(v)
  403. s = _check_func(check_len, f'length >= {min_len} and length <= {max_len}', s)
  404. elif isinstance(constraint, annotated_types.MultipleOf):
  405. multiple_of = constraint.multiple_of
  406. if s and s['type'] in {'int', 'float', 'decimal'}:
  407. s = s.copy()
  408. if s['type'] == 'int' and isinstance(multiple_of, int):
  409. s['multiple_of'] = multiple_of
  410. elif s['type'] == 'float' and isinstance(multiple_of, float):
  411. s['multiple_of'] = multiple_of
  412. elif s['type'] == 'decimal' and isinstance(multiple_of, Decimal):
  413. s['multiple_of'] = multiple_of
  414. def check_multiple_of(v: Any) -> bool:
  415. return v % multiple_of == 0
  416. s = _check_func(check_multiple_of, f'% {multiple_of} == 0', s)
  417. elif isinstance(constraint, annotated_types.Timezone):
  418. tz = constraint.tz
  419. if tz is ...:
  420. if s and s['type'] == 'datetime':
  421. s = s.copy()
  422. s['tz_constraint'] = 'aware'
  423. else:
  424. def check_tz_aware(v: object) -> bool:
  425. assert isinstance(v, datetime.datetime)
  426. return v.tzinfo is not None
  427. s = _check_func(check_tz_aware, 'timezone aware', s)
  428. elif tz is None:
  429. if s and s['type'] == 'datetime':
  430. s = s.copy()
  431. s['tz_constraint'] = 'naive'
  432. else:
  433. def check_tz_naive(v: object) -> bool:
  434. assert isinstance(v, datetime.datetime)
  435. return v.tzinfo is None
  436. s = _check_func(check_tz_naive, 'timezone naive', s)
  437. else:
  438. raise NotImplementedError('Constraining to a specific timezone is not yet supported')
  439. elif isinstance(constraint, annotated_types.Interval):
  440. if constraint.ge:
  441. s = _apply_constraint(s, annotated_types.Ge(constraint.ge))
  442. if constraint.gt:
  443. s = _apply_constraint(s, annotated_types.Gt(constraint.gt))
  444. if constraint.le:
  445. s = _apply_constraint(s, annotated_types.Le(constraint.le))
  446. if constraint.lt:
  447. s = _apply_constraint(s, annotated_types.Lt(constraint.lt))
  448. assert s is not None
  449. elif isinstance(constraint, annotated_types.Predicate):
  450. func = constraint.func
  451. # Same logic as in `_known_annotated_metadata.apply_known_metadata()`:
  452. predicate_name = f'{func.__qualname__!r} ' if hasattr(func, '__qualname__') else ''
  453. def predicate_func(v: Any) -> Any:
  454. if not func(v):
  455. raise PydanticCustomError(
  456. 'predicate_failed',
  457. f'Predicate {predicate_name}failed', # pyright: ignore[reportArgumentType]
  458. )
  459. return v
  460. if s is None:
  461. s = cs.no_info_plain_validator_function(predicate_func)
  462. else:
  463. s = cs.no_info_after_validator_function(predicate_func, s)
  464. elif isinstance(constraint, _NotEq):
  465. value = constraint.value
  466. def check_not_eq(v: Any) -> bool:
  467. return operator.__ne__(v, value)
  468. s = _check_func(check_not_eq, f'!= {value}', s)
  469. elif isinstance(constraint, _Eq):
  470. value = constraint.value
  471. def check_eq(v: Any) -> bool:
  472. return operator.__eq__(v, value)
  473. s = _check_func(check_eq, f'== {value}', s)
  474. elif isinstance(constraint, _In):
  475. values = constraint.values
  476. def check_in(v: Any) -> bool:
  477. return operator.__contains__(values, v)
  478. s = _check_func(check_in, f'in {values}', s)
  479. elif isinstance(constraint, _NotIn):
  480. values = constraint.values
  481. def check_not_in(v: Any) -> bool:
  482. return operator.__not__(operator.__contains__(values, v))
  483. s = _check_func(check_not_in, f'not in {values}', s)
  484. else:
  485. assert isinstance(constraint, Pattern)
  486. if s and s['type'] == 'str':
  487. s = s.copy()
  488. s['pattern'] = constraint.pattern
  489. else:
  490. def check_pattern(v: object) -> bool:
  491. assert isinstance(v, str)
  492. return constraint.match(v) is not None
  493. s = _check_func(check_pattern, f'~ {constraint.pattern}', s)
  494. return s
  495. class _SupportsRange(annotated_types.SupportsLe, annotated_types.SupportsGe, Protocol):
  496. pass
  497. class _SupportsLen(Protocol):
  498. def __len__(self) -> int: ...
  499. _NewOutGt = TypeVar('_NewOutGt', bound=annotated_types.SupportsGt)
  500. _NewOutGe = TypeVar('_NewOutGe', bound=annotated_types.SupportsGe)
  501. _NewOutLt = TypeVar('_NewOutLt', bound=annotated_types.SupportsLt)
  502. _NewOutLe = TypeVar('_NewOutLe', bound=annotated_types.SupportsLe)
  503. _NewOutLen = TypeVar('_NewOutLen', bound=_SupportsLen)
  504. _NewOutDiv = TypeVar('_NewOutDiv', bound=annotated_types.SupportsDiv)
  505. _NewOutMod = TypeVar('_NewOutMod', bound=annotated_types.SupportsMod)
  506. _NewOutDatetime = TypeVar('_NewOutDatetime', bound=datetime.datetime)
  507. _NewOutInterval = TypeVar('_NewOutInterval', bound=_SupportsRange)
  508. _OtherIn = TypeVar('_OtherIn')
  509. _OtherOut = TypeVar('_OtherOut')