1# This file is part of Hypothesis, which may be found at
2# https://github.com/HypothesisWorks/hypothesis/
3#
4# Copyright the Hypothesis Authors.
5# Individual contributors are listed in AUTHORS.rst and the git log.
6#
7# This Source Code Form is subject to the terms of the Mozilla Public License,
8# v. 2.0. If a copy of the MPL was not distributed with this file, You can
9# obtain one at https://mozilla.org/MPL/2.0/.
10
11import ast
12import dis
13import hashlib
14import inspect
15import linecache
16import sys
17import textwrap
18from collections.abc import Callable, MutableMapping
19from inspect import Parameter
20from typing import Any
21from weakref import WeakKeyDictionary
22
23from hypothesis.internal import reflection
24from hypothesis.internal.cache import LRUCache
25
26# we have several levels of caching for lambda descriptions.
27# * LAMBDA_DESCRIPTION_CACHE maps a lambda f to its description _lambda_description(f).
28# Note that _lambda_description(f) may not be identical to f as it appears in the
29# source code file.
30# * LAMBDA_DIGEST_DESCRIPTION_CACHE maps _function_key(f) to _lambda_description(f).
31# _function_key implements something close to "ast equality":
32# two syntactically identical (minus whitespace etc) lambdas appearing in
33# different files have the same key. Cache hits here provide a fast path which
34# avoids ast-parsing syntactic lambdas we've seen before. Two lambdas with the
35# same _function_key will not have different _lambda_descriptions - if
36# they do, that's a bug here.
37# * AST_LAMBDAS_CACHE maps source code lines to a list of the lambdas found in
38# that source code. A cache hit here avoids reparsing the ast.
39LAMBDA_DESCRIPTION_CACHE: MutableMapping[Callable, str] = WeakKeyDictionary()
40LAMBDA_DIGEST_DESCRIPTION_CACHE: LRUCache[tuple[Any], str] = LRUCache(max_size=1000)
41AST_LAMBDAS_CACHE: LRUCache[tuple[str], list[ast.Lambda]] = LRUCache(max_size=100)
42
43
44def extract_all_lambdas(tree):
45 lambdas = []
46
47 class Visitor(ast.NodeVisitor):
48
49 def visit_Lambda(self, node):
50 lambdas.append(node)
51 self.visit(node.body)
52
53 Visitor().visit(tree)
54 return lambdas
55
56
57def extract_all_attributes(tree):
58 attributes = []
59
60 class Visitor(ast.NodeVisitor):
61 def visit_Attribute(self, node):
62 attributes.append(node)
63 self.visit(node.value)
64
65 Visitor().visit(tree)
66 return attributes
67
68
69def _function_key(f, *, bounded_size=False, ignore_name=False):
70 """Returns a digest that differentiates functions that have different sources.
71
72 Either a function or a code object may be passed. If code object, default
73 arg/kwarg values are not recoverable - this is the best we can do, and is
74 sufficient for the use case of comparing nested lambdas.
75 """
76 try:
77 code = f.__code__
78 defaults_repr = repr((f.__defaults__, f.__kwdefaults__))
79 except AttributeError:
80 code = f
81 defaults_repr = ()
82 consts_repr = repr(code.co_consts)
83 if bounded_size:
84 # Compress repr to avoid keeping arbitrarily large strings pinned as cache
85 # keys. We don't do this unconditionally because hashing takes time, and is
86 # not necessary if the key is used just for comparison (and is not stored).
87 if len(consts_repr) > 48:
88 consts_repr = hashlib.sha384(consts_repr.encode()).digest()
89 if len(defaults_repr) > 48:
90 defaults_repr = hashlib.sha384(defaults_repr.encode()).digest()
91 return (
92 consts_repr,
93 defaults_repr,
94 code.co_argcount,
95 code.co_kwonlyargcount,
96 code.co_code,
97 code.co_names,
98 code.co_varnames,
99 code.co_freevars,
100 ignore_name or code.co_name,
101 )
102
103
104class _op:
105 # Look up the opcode values dynamically, since they can change between versions. If
106 # the opcode does not exist on a version, we set it to None.
107 NOP = dis.opmap["NOP"]
108 LOAD_FAST = dis.opmap["LOAD_FAST"]
109 LOAD_FAST_LOAD_FAST = (
110 dis.opmap["LOAD_FAST_LOAD_FAST"] if sys.version_info[:2] >= (3, 13) else None
111 )
112 LOAD_FAST_BORROW = (
113 dis.opmap["LOAD_FAST_BORROW"] if sys.version_info[:2] >= (3, 14) else None
114 )
115 LOAD_FAST_BORROW_LOAD_FAST_BORROW = (
116 dis.opmap["LOAD_FAST_BORROW_LOAD_FAST_BORROW"]
117 if sys.version_info[:2] >= (3, 14)
118 else None
119 )
120
121
122def _normalize_code(f, l):
123 # A small selection of possible peephole code transformations, based on what
124 # is actually seen to differ between compilations in our test suite. Each
125 # entry contains two equivalent opcode sequences, plus a condition
126 # function called with their respective oparg sequences, which must return
127 # true for the transformation to be valid.
128 Checker = Callable[[list[int], list[int]], bool]
129 transforms: tuple[list[int], list[int], Checker | None] = [
130 ([_op.NOP], [], lambda a, b: True),
131 (
132 [_op.LOAD_FAST, _op.LOAD_FAST],
133 [_op.LOAD_FAST_LOAD_FAST],
134 lambda a, b: a == [b[0] >> 4, b[0] & 15],
135 ),
136 (
137 [_op.LOAD_FAST_BORROW, _op.LOAD_FAST_BORROW],
138 [_op.LOAD_FAST_BORROW_LOAD_FAST_BORROW],
139 lambda a, b: a == [b[0] >> 4, b[0] & 15],
140 ),
141 ]
142 # Avoid applying any transform with an opcode that doesn't exist on this python version.
143 transforms = [t for t in transforms if None not in t[0] + t[1]]
144 # augment with converse
145 transforms += [
146 (
147 ops_b,
148 ops_a,
149 condition and (lambda a, b, condition=condition: condition(b, a)),
150 )
151 for ops_a, ops_b, condition in transforms
152 ]
153
154 # Normalize equivalent code. We assume that each bytecode op is 2 bytes,
155 # which is the case since Python 3.6. Since the opcodes values may change
156 # between version, there is a risk that a transform may not be equivalent
157 # -- even so, the risk of a bad transform producing a false positive is
158 # minuscule.
159 co_code = list(l.__code__.co_code)
160 f_code = list(f.__code__.co_code)
161
162 def alternating(code, i, n):
163 return code[i : i + 2 * n : 2]
164
165 i = 2
166 while i < max(len(co_code), len(f_code)):
167 # note that co_code is mutated in loop
168 if i < min(len(co_code), len(f_code)) and f_code[i] == co_code[i]:
169 i += 2
170 else:
171 for op1, op2, condition in transforms:
172 if (
173 op1 == alternating(f_code, i, len(op1))
174 and op2 == alternating(co_code, i, len(op2))
175 and condition(
176 alternating(f_code, i + 1, len(op1)),
177 alternating(co_code, i + 1, len(op2)),
178 )
179 ):
180 break
181 else:
182 # no point in continuing since the bytecodes are different anyway
183 break
184 # Splice in the transform and continue
185 co_code = (
186 co_code[:i] + f_code[i : i + 2 * len(op1)] + co_code[i + 2 * len(op2) :]
187 )
188 i += 2 * len(op1)
189
190 # Normalize consts, in particular replace any lambda consts with the
191 # corresponding const from the template function, IFF they have the same
192 # source key.
193
194 f_consts = f.__code__.co_consts
195 l_consts = l.__code__.co_consts
196 if len(f_consts) == len(l_consts) and any(
197 inspect.iscode(l_const) for l_const in l_consts
198 ):
199 normalized_consts = []
200 for f_const, l_const in zip(f_consts, l_consts, strict=True):
201 if (
202 inspect.iscode(l_const)
203 and inspect.iscode(f_const)
204 and _function_key(f_const) == _function_key(l_const)
205 ):
206 # If the lambdas are compiled from the same source, make them be the
207 # same object so that the toplevel lambdas end up equal. Note that
208 # default arguments are not available on the code objects. But if the
209 # default arguments differ then the lambdas must also differ in other
210 # ways, since default arguments are set up from bytecode and constants.
211 # I.e., this appears to be safe wrt false positives.
212 normalized_consts.append(f_const)
213 else:
214 normalized_consts.append(l_const)
215 else:
216 normalized_consts = l_consts
217
218 return l.__code__.replace(
219 co_code=bytes(co_code),
220 co_consts=tuple(normalized_consts),
221 )
222
223
224_module_map: dict[int, str] = {}
225
226
227def _mimic_lambda_from_node(f, node):
228 # Compile the source (represented by an ast.Lambda node) in a context that
229 # as far as possible mimics the context that f was compiled in. If - and
230 # only if - this was the source of f then the result is indistinguishable
231 # from f itself (to a casual observer such as _function_key).
232 f_globals = f.__globals__.copy()
233 f_code = f.__code__
234 source = ast.unparse(node)
235
236 # Install values for non-literal argument defaults. Thankfully, these are
237 # always captured by value - so there is no interaction with the closure.
238 if f.__defaults__:
239 for f_default, l_default in zip(
240 f.__defaults__, node.args.defaults, strict=True
241 ):
242 if isinstance(l_default, ast.Name):
243 f_globals[l_default.id] = f_default
244 if f.__kwdefaults__: # pragma: no cover
245 for l_default, l_varname in zip(
246 node.args.kw_defaults, node.args.kwonlyargs, strict=True
247 ):
248 if isinstance(l_default, ast.Name):
249 f_globals[l_default.id] = f.__kwdefaults__[l_varname.arg]
250
251 # CPython's compiler treats known imports differently than normal globals,
252 # so check if we use attributes from globals that are modules (if so, we
253 # import them explicitly and redundantly in the exec below)
254 referenced_modules = [
255 (local_name, module)
256 for attr in extract_all_attributes(node)
257 if (
258 isinstance(attr.value, ast.Name)
259 and (local_name := attr.value.id)
260 and inspect.ismodule(module := f_globals.get(local_name))
261 )
262 ]
263
264 if not f_code.co_freevars and not referenced_modules:
265 compiled = eval(source, f_globals)
266 else:
267 if f_code.co_freevars:
268 # We have to reconstruct a local closure. The closure will have
269 # the same values as the original function, although this is not
270 # required for source/bytecode equality.
271 f_globals |= {
272 f"__lc{i}": c.cell_contents for i, c in enumerate(f.__closure__)
273 }
274 captures = [f"{name}=__lc{i}" for i, name in enumerate(f_code.co_freevars)]
275 capture_str = ";".join(captures) + ";"
276 else:
277 capture_str = ""
278 if referenced_modules:
279 # We add import statements for all referenced modules, since that
280 # influences the compiled code. The assumption is that these modules
281 # were explicitly imported, not assigned, in the source - if not,
282 # this may/will give a different compilation result.
283 global _module_map
284 if len(_module_map) != len(sys.modules): # pragma: no branch
285 _module_map = {id(module): name for name, module in sys.modules.items()}
286 imports = [
287 (module_name, local_name)
288 for local_name, module in referenced_modules
289 if (module_name := _module_map.get(id(module))) is not None
290 ]
291 import_fragments = [f"{name} as {asname}" for name, asname in set(imports)]
292 import_str = f"import {','.join(import_fragments)}\n"
293 else:
294 import_str = ""
295 exec_str = (
296 f"{import_str}def __construct_lambda(): {capture_str} return ({source})"
297 )
298 exec(exec_str, f_globals)
299 compiled = f_globals["__construct_lambda"]()
300
301 return compiled
302
303
304def _lambda_code_matches_node(f, node):
305 try:
306 compiled = _mimic_lambda_from_node(f, node)
307 except (NameError, SyntaxError): # pragma: no cover # source is generated from ast
308 return False
309 if _function_key(f) == _function_key(compiled):
310 return True
311 # Try harder
312 compiled.__code__ = _normalize_code(f, compiled)
313 return _function_key(f) == _function_key(compiled)
314
315
316def _check_unknown_perfectly_aligned_lambda(candidate): # pragma: no cover
317 # This is a monkeypatch point for our self-tests, to make unknown
318 # lambdas raise.
319 pass
320
321
322def _lambda_description(f, leeway=50, *, fail_if_confused_with_perfect_candidate=False):
323 if hasattr(f, "__wrapped_target"):
324 f = f.__wrapped_target
325
326 # You might be wondering how a lambda can have a return-type annotation?
327 # The answer is that we add this at runtime, in new_given_signature(),
328 # and we do support strange choices as applying @given() to a lambda.
329 sig = inspect.signature(f)
330 assert sig.return_annotation in (Parameter.empty, None), sig
331
332 # Using pytest-xdist on Python 3.13, there's an entry in the linecache for
333 # file "<string>", which then returns nonsense to getsource. Discard it.
334 linecache.cache.pop("<string>", None)
335
336 def format_lambda(body):
337 # The signature is more informative than the corresponding ast.unparse
338 # output in the case of default argument values, so add the signature
339 # to the unparsed body
340 return (
341 f"lambda {str(sig)[1:-1]}: {body}" if sig.parameters else f"lambda: {body}"
342 )
343
344 if_confused = format_lambda("<unknown>")
345
346 try:
347 source_lines, lineno0 = inspect.findsource(f)
348 source_lines = tuple(source_lines) # make it hashable
349 except OSError:
350 return if_confused
351
352 try:
353 all_lambdas = AST_LAMBDAS_CACHE[source_lines]
354 except KeyError:
355 # The source isn't already parsed, so we try to shortcut by parsing just
356 # the local block. If that fails to produce a code-identical lambda,
357 # fall through to the full parse.
358 local_lines = inspect.getblock(source_lines[lineno0:])
359 local_block = textwrap.dedent("".join(local_lines))
360 # The fairly common ".map(lambda x: ...)" case. This partial block
361 # isn't valid syntax, but it might be if we remove the leading ".".
362 local_block = local_block.removeprefix(".")
363
364 try:
365 local_tree = ast.parse(local_block)
366 except SyntaxError:
367 pass
368 else:
369 local_lambdas = extract_all_lambdas(local_tree)
370 for candidate in local_lambdas:
371 if reflection.ast_arguments_matches_signature(
372 candidate.args, sig
373 ) and _lambda_code_matches_node(f, candidate):
374 return format_lambda(ast.unparse(candidate.body))
375
376 # Local parse failed or didn't produce a match, go ahead with the full parse
377 try:
378 tree = ast.parse("".join(source_lines))
379 except SyntaxError:
380 all_lambdas = []
381 else:
382 all_lambdas = extract_all_lambdas(tree)
383 AST_LAMBDAS_CACHE[source_lines] = all_lambdas
384
385 aligned_lambdas = []
386 for candidate in all_lambdas:
387 if (
388 candidate.lineno - leeway <= lineno0 + 1 <= candidate.lineno + leeway
389 and reflection.ast_arguments_matches_signature(candidate.args, sig)
390 ):
391 aligned_lambdas.append(candidate)
392
393 aligned_lambdas.sort(key=lambda c: abs(lineno0 + 1 - c.lineno))
394 for candidate in aligned_lambdas:
395 if _lambda_code_matches_node(f, candidate):
396 return format_lambda(ast.unparse(candidate.body))
397
398 # None of the aligned lambdas match perfectly in generated code.
399 if aligned_lambdas and aligned_lambdas[0].lineno == lineno0 + 1:
400 _check_unknown_perfectly_aligned_lambda(aligned_lambdas[0])
401
402 return if_confused
403
404
405def lambda_description(f):
406 """
407 Returns a syntactically-valid expression describing `f`. This is often, but
408 not always, the exact lambda definition string which appears in the source code.
409 The difference comes from parsing the lambda ast into `tree` and then returning
410 the result of `ast.unparse(tree)`, which may differ in whitespace, double vs
411 single quotes, etc.
412
413 Returns a string indicating an unknown body if the parsing gets confused in any way.
414 """
415 try:
416 return LAMBDA_DESCRIPTION_CACHE[f]
417 except KeyError:
418 pass
419
420 # Follow the breadcrumb left by impersonate() before computing the cache key,
421 # because every wrapper we generate for a given signature has identical code
422 # and so would otherwise share a key with unrelated wrapped functions.
423 target = getattr(f, "__wrapped_target", f)
424 key = _function_key(target, bounded_size=True)
425 location = (target.__code__.co_filename, target.__code__.co_firstlineno)
426 try:
427 description, failed_locations = LAMBDA_DIGEST_DESCRIPTION_CACHE[key]
428 except KeyError:
429 failed_locations = set()
430 else:
431 # We got a hit in the digests cache, but only use it if either it has
432 # a good (known) description, or if it is unknown but we already tried
433 # to parse its exact source location before.
434 if "<unknown>" not in description or location in failed_locations:
435 # use the cached result
436 LAMBDA_DESCRIPTION_CACHE[f] = description
437 return description
438
439 description = _lambda_description(f)
440 LAMBDA_DESCRIPTION_CACHE[f] = description
441 if "<unknown>" in description:
442 failed_locations.add(location)
443 else:
444 failed_locations.clear() # we have a good description now
445 LAMBDA_DIGEST_DESCRIPTION_CACHE[key] = description, failed_locations
446 return description