Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.
 
 
 
 

365 rindas
13 KiB

  1. """
  2. Register Hypothesis strategies for Pydantic custom types.
  3. This enables fully-automatic generation of test data for most Pydantic classes.
  4. Note that this module has *no* runtime impact on Pydantic itself; instead it
  5. is registered as a setuptools entry point and Hypothesis will import it if
  6. Pydantic is installed. See also:
  7. https://hypothesis.readthedocs.io/en/latest/strategies.html#registering-strategies-via-setuptools-entry-points
  8. https://hypothesis.readthedocs.io/en/latest/data.html#hypothesis.strategies.register_type_strategy
  9. https://hypothesis.readthedocs.io/en/latest/strategies.html#interaction-with-pytest-cov
  10. https://pydantic-docs.helpmanual.io/usage/types/#pydantic-types
  11. Note that because our motivation is to *improve user experience*, the strategies
  12. are always sound (never generate invalid data) but sacrifice completeness for
  13. maintainability (ie may be unable to generate some tricky but valid data).
  14. Finally, this module makes liberal use of `# type: ignore[<code>]` pragmas.
  15. This is because Hypothesis annotates `register_type_strategy()` with
  16. `(T, SearchStrategy[T])`, but in most cases we register e.g. `ConstrainedInt`
  17. to generate instances of the builtin `int` type which match the constraints.
  18. """
  19. import contextlib
  20. import ipaddress
  21. import json
  22. import math
  23. from fractions import Fraction
  24. from typing import Callable, Dict, Type, Union, cast, overload
  25. import hypothesis.strategies as st
  26. import pydantic
  27. import pydantic.color
  28. import pydantic.types
  29. # FilePath and DirectoryPath are explicitly unsupported, as we'd have to create
  30. # them on-disk, and that's unsafe in general without being told *where* to do so.
  31. #
  32. # URLs are unsupported because it's easy for users to define their own strategy for
  33. # "normal" URLs, and hard for us to define a general strategy which includes "weird"
  34. # URLs but doesn't also have unpredictable performance problems.
  35. #
  36. # conlist() and conset() are unsupported for now, because the workarounds for
  37. # Cython and Hypothesis to handle parametrized generic types are incompatible.
  38. # Once Cython can support 'normal' generics we'll revisit this.
  39. # Emails
  40. try:
  41. import email_validator
  42. except ImportError: # pragma: no cover
  43. pass
  44. else:
  45. def is_valid_email(s: str) -> bool:
  46. # Hypothesis' st.emails() occasionally generates emails like 0@A0--0.ac
  47. # that are invalid according to email-validator, so we filter those out.
  48. try:
  49. email_validator.validate_email(s, check_deliverability=False)
  50. return True
  51. except email_validator.EmailNotValidError: # pragma: no cover
  52. return False
  53. # Note that these strategies deliberately stay away from any tricky Unicode
  54. # or other encoding issues; we're just trying to generate *something* valid.
  55. st.register_type_strategy(pydantic.EmailStr, st.emails().filter(is_valid_email)) # type: ignore[arg-type]
  56. st.register_type_strategy(
  57. pydantic.NameEmail,
  58. st.builds(
  59. '{} <{}>'.format, # type: ignore[arg-type]
  60. st.from_regex('[A-Za-z0-9_]+( [A-Za-z0-9_]+){0,5}', fullmatch=True),
  61. st.emails().filter(is_valid_email),
  62. ),
  63. )
  64. # PyObject - dotted names, in this case taken from the math module.
  65. st.register_type_strategy(
  66. pydantic.PyObject, # type: ignore[arg-type]
  67. st.sampled_from(
  68. [cast(pydantic.PyObject, f'math.{name}') for name in sorted(vars(math)) if not name.startswith('_')]
  69. ),
  70. )
  71. # CSS3 Colors; as name, hex, rgb(a) tuples or strings, or hsl strings
  72. _color_regexes = (
  73. '|'.join(
  74. (
  75. pydantic.color.r_hex_short,
  76. pydantic.color.r_hex_long,
  77. pydantic.color.r_rgb,
  78. pydantic.color.r_rgba,
  79. pydantic.color.r_hsl,
  80. pydantic.color.r_hsla,
  81. )
  82. )
  83. # Use more precise regex patterns to avoid value-out-of-range errors
  84. .replace(pydantic.color._r_sl, r'(?:(\d\d?(?:\.\d+)?|100(?:\.0+)?)%)')
  85. .replace(pydantic.color._r_alpha, r'(?:(0(?:\.\d+)?|1(?:\.0+)?|\.\d+|\d{1,2}%))')
  86. .replace(pydantic.color._r_255, r'(?:((?:\d|\d\d|[01]\d\d|2[0-4]\d|25[0-4])(?:\.\d+)?|255(?:\.0+)?))')
  87. )
  88. st.register_type_strategy(
  89. pydantic.color.Color,
  90. st.one_of(
  91. st.sampled_from(sorted(pydantic.color.COLORS_BY_NAME)),
  92. st.tuples(
  93. st.integers(0, 255),
  94. st.integers(0, 255),
  95. st.integers(0, 255),
  96. st.none() | st.floats(0, 1) | st.floats(0, 100).map('{}%'.format),
  97. ),
  98. st.from_regex(_color_regexes, fullmatch=True),
  99. ),
  100. )
  101. # Card numbers, valid according to the Luhn algorithm
  102. def add_luhn_digit(card_number: str) -> str:
  103. # See https://en.wikipedia.org/wiki/Luhn_algorithm
  104. for digit in '0123456789':
  105. with contextlib.suppress(Exception):
  106. pydantic.PaymentCardNumber.validate_luhn_check_digit(card_number + digit)
  107. return card_number + digit
  108. raise AssertionError('Unreachable') # pragma: no cover
  109. card_patterns = (
  110. # Note that these patterns omit the Luhn check digit; that's added by the function above
  111. '4[0-9]{14}', # Visa
  112. '5[12345][0-9]{13}', # Mastercard
  113. '3[47][0-9]{12}', # American Express
  114. '[0-26-9][0-9]{10,17}', # other (incomplete to avoid overlap)
  115. )
  116. st.register_type_strategy(
  117. pydantic.PaymentCardNumber,
  118. st.from_regex('|'.join(card_patterns), fullmatch=True).map(add_luhn_digit), # type: ignore[arg-type]
  119. )
  120. # UUIDs
  121. st.register_type_strategy(pydantic.UUID1, st.uuids(version=1))
  122. st.register_type_strategy(pydantic.UUID3, st.uuids(version=3))
  123. st.register_type_strategy(pydantic.UUID4, st.uuids(version=4))
  124. st.register_type_strategy(pydantic.UUID5, st.uuids(version=5))
  125. # Secrets
  126. st.register_type_strategy(pydantic.SecretBytes, st.binary().map(pydantic.SecretBytes))
  127. st.register_type_strategy(pydantic.SecretStr, st.text().map(pydantic.SecretStr))
  128. # IP addresses, networks, and interfaces
  129. st.register_type_strategy(pydantic.IPvAnyAddress, st.ip_addresses()) # type: ignore[arg-type]
  130. st.register_type_strategy(
  131. pydantic.IPvAnyInterface,
  132. st.from_type(ipaddress.IPv4Interface) | st.from_type(ipaddress.IPv6Interface), # type: ignore[arg-type]
  133. )
  134. st.register_type_strategy(
  135. pydantic.IPvAnyNetwork,
  136. st.from_type(ipaddress.IPv4Network) | st.from_type(ipaddress.IPv6Network), # type: ignore[arg-type]
  137. )
  138. # We hook into the con***() functions and the ConstrainedNumberMeta metaclass,
  139. # so here we only have to register subclasses for other constrained types which
  140. # don't go via those mechanisms. Then there are the registration hooks below.
  141. st.register_type_strategy(pydantic.StrictBool, st.booleans())
  142. st.register_type_strategy(pydantic.StrictStr, st.text())
  143. # Constrained-type resolver functions
  144. #
  145. # For these ones, we actually want to inspect the type in order to work out a
  146. # satisfying strategy. First up, the machinery for tracking resolver functions:
  147. RESOLVERS: Dict[type, Callable[[type], st.SearchStrategy]] = {} # type: ignore[type-arg]
  148. @overload
  149. def _registered(typ: Type[pydantic.types.T]) -> Type[pydantic.types.T]:
  150. pass
  151. @overload
  152. def _registered(typ: pydantic.types.ConstrainedNumberMeta) -> pydantic.types.ConstrainedNumberMeta:
  153. pass
  154. def _registered(
  155. typ: Union[Type[pydantic.types.T], pydantic.types.ConstrainedNumberMeta]
  156. ) -> Union[Type[pydantic.types.T], pydantic.types.ConstrainedNumberMeta]:
  157. # This function replaces the version in `pydantic.types`, in order to
  158. # effect the registration of new constrained types so that Hypothesis
  159. # can generate valid examples.
  160. pydantic.types._DEFINED_TYPES.add(typ)
  161. for supertype, resolver in RESOLVERS.items():
  162. if issubclass(typ, supertype):
  163. st.register_type_strategy(typ, resolver(typ)) # type: ignore
  164. return typ
  165. raise NotImplementedError(f'Unknown type {typ!r} has no resolver to register') # pragma: no cover
  166. def resolves(
  167. typ: Union[type, pydantic.types.ConstrainedNumberMeta]
  168. ) -> Callable[[Callable[..., st.SearchStrategy]], Callable[..., st.SearchStrategy]]: # type: ignore[type-arg]
  169. def inner(f): # type: ignore
  170. assert f not in RESOLVERS
  171. RESOLVERS[typ] = f
  172. return f
  173. return inner
  174. # Type-to-strategy resolver functions
  175. @resolves(pydantic.JsonWrapper)
  176. def resolve_json(cls): # type: ignore[no-untyped-def]
  177. try:
  178. inner = st.none() if cls.inner_type is None else st.from_type(cls.inner_type)
  179. except Exception: # pragma: no cover
  180. finite = st.floats(allow_infinity=False, allow_nan=False)
  181. inner = st.recursive(
  182. base=st.one_of(st.none(), st.booleans(), st.integers(), finite, st.text()),
  183. extend=lambda x: st.lists(x) | st.dictionaries(st.text(), x), # type: ignore
  184. )
  185. return st.builds(
  186. json.dumps,
  187. inner,
  188. ensure_ascii=st.booleans(),
  189. indent=st.none() | st.integers(0, 16),
  190. sort_keys=st.booleans(),
  191. )
  192. @resolves(pydantic.ConstrainedBytes)
  193. def resolve_conbytes(cls): # type: ignore[no-untyped-def] # pragma: no cover
  194. min_size = cls.min_length or 0
  195. max_size = cls.max_length
  196. if not cls.strip_whitespace:
  197. return st.binary(min_size=min_size, max_size=max_size)
  198. # Fun with regex to ensure we neither start nor end with whitespace
  199. repeats = '{{{},{}}}'.format(
  200. min_size - 2 if min_size > 2 else 0,
  201. max_size - 2 if (max_size or 0) > 2 else '',
  202. )
  203. if min_size >= 2:
  204. pattern = rf'\W.{repeats}\W'
  205. elif min_size == 1:
  206. pattern = rf'\W(.{repeats}\W)?'
  207. else:
  208. assert min_size == 0
  209. pattern = rf'(\W(.{repeats}\W)?)?'
  210. return st.from_regex(pattern.encode(), fullmatch=True)
  211. @resolves(pydantic.ConstrainedDecimal)
  212. def resolve_condecimal(cls): # type: ignore[no-untyped-def]
  213. min_value = cls.ge
  214. max_value = cls.le
  215. if cls.gt is not None:
  216. assert min_value is None, 'Set `gt` or `ge`, but not both'
  217. min_value = cls.gt
  218. if cls.lt is not None:
  219. assert max_value is None, 'Set `lt` or `le`, but not both'
  220. max_value = cls.lt
  221. s = st.decimals(min_value, max_value, allow_nan=False, places=cls.decimal_places)
  222. if cls.lt is not None:
  223. s = s.filter(lambda d: d < cls.lt)
  224. if cls.gt is not None:
  225. s = s.filter(lambda d: cls.gt < d)
  226. return s
  227. @resolves(pydantic.ConstrainedFloat)
  228. def resolve_confloat(cls): # type: ignore[no-untyped-def]
  229. min_value = cls.ge
  230. max_value = cls.le
  231. exclude_min = False
  232. exclude_max = False
  233. if cls.gt is not None:
  234. assert min_value is None, 'Set `gt` or `ge`, but not both'
  235. min_value = cls.gt
  236. exclude_min = True
  237. if cls.lt is not None:
  238. assert max_value is None, 'Set `lt` or `le`, but not both'
  239. max_value = cls.lt
  240. exclude_max = True
  241. if cls.multiple_of is None:
  242. return st.floats(min_value, max_value, exclude_min=exclude_min, exclude_max=exclude_max, allow_nan=False)
  243. if min_value is not None:
  244. min_value = math.ceil(min_value / cls.multiple_of)
  245. if exclude_min:
  246. min_value = min_value + 1
  247. if max_value is not None:
  248. assert max_value >= cls.multiple_of, 'Cannot build model with max value smaller than multiple of'
  249. max_value = math.floor(max_value / cls.multiple_of)
  250. if exclude_max:
  251. max_value = max_value - 1
  252. return st.integers(min_value, max_value).map(lambda x: x * cls.multiple_of)
  253. @resolves(pydantic.ConstrainedInt)
  254. def resolve_conint(cls): # type: ignore[no-untyped-def]
  255. min_value = cls.ge
  256. max_value = cls.le
  257. if cls.gt is not None:
  258. assert min_value is None, 'Set `gt` or `ge`, but not both'
  259. min_value = cls.gt + 1
  260. if cls.lt is not None:
  261. assert max_value is None, 'Set `lt` or `le`, but not both'
  262. max_value = cls.lt - 1
  263. if cls.multiple_of is None or cls.multiple_of == 1:
  264. return st.integers(min_value, max_value)
  265. # These adjustments and the .map handle integer-valued multiples, while the
  266. # .filter handles trickier cases as for confloat.
  267. if min_value is not None:
  268. min_value = math.ceil(Fraction(min_value) / Fraction(cls.multiple_of))
  269. if max_value is not None:
  270. max_value = math.floor(Fraction(max_value) / Fraction(cls.multiple_of))
  271. return st.integers(min_value, max_value).map(lambda x: x * cls.multiple_of)
  272. @resolves(pydantic.ConstrainedStr)
  273. def resolve_constr(cls): # type: ignore[no-untyped-def] # pragma: no cover
  274. min_size = cls.min_length or 0
  275. max_size = cls.max_length
  276. if cls.regex is None and not cls.strip_whitespace:
  277. return st.text(min_size=min_size, max_size=max_size)
  278. if cls.regex is not None:
  279. strategy = st.from_regex(cls.regex)
  280. if cls.strip_whitespace:
  281. strategy = strategy.filter(lambda s: s == s.strip())
  282. elif cls.strip_whitespace:
  283. repeats = '{{{},{}}}'.format(
  284. min_size - 2 if min_size > 2 else 0,
  285. max_size - 2 if (max_size or 0) > 2 else '',
  286. )
  287. if min_size >= 2:
  288. strategy = st.from_regex(rf'\W.{repeats}\W')
  289. elif min_size == 1:
  290. strategy = st.from_regex(rf'\W(.{repeats}\W)?')
  291. else:
  292. assert min_size == 0
  293. strategy = st.from_regex(rf'(\W(.{repeats}\W)?)?')
  294. if min_size == 0 and max_size is None:
  295. return strategy
  296. elif max_size is None:
  297. return strategy.filter(lambda s: min_size <= len(s))
  298. return strategy.filter(lambda s: min_size <= len(s) <= max_size)
  299. # Finally, register all previously-defined types, and patch in our new function
  300. for typ in pydantic.types._DEFINED_TYPES:
  301. _registered(typ)
  302. pydantic.types._registered = _registered
  303. st.register_type_strategy(pydantic.Json, resolve_json)