Repository: HypothesisWorks/hypothesis

Issue:
Ghostwriter: try to import required parts for custom strategies
As noted in @rsokl's [`ghosty` demo project](https://github.com/rsokl/ghosty), ghostwritten tests sometimes fail with a `NameError`, ultimately due to a missing import `Foo` in `builds(Foo)`, or `f` in methods like `.filter(f)`, `.map(f)`, or `.flatmap(f)`.

We would obviously like to fix this - personally I expect that it's intractable in general, and would therefore be happy with a shallow and/or heuristic solution... though I certainly wouldn't turn down a complete fix if one is available.


Discussion:

Resolving patch:
diff --git a/hypothesis-python/src/hypothesis/extra/ghostwriter.py b/hypothesis-python/src/hypothesis/extra/ghostwriter.py
--- a/hypothesis-python/src/hypothesis/extra/ghostwriter.py
+++ b/hypothesis-python/src/hypothesis/extra/ghostwriter.py
@@ -94,11 +94,18 @@
 import black
 
 from hypothesis import find, strategies as st
-from hypothesis.errors import InvalidArgument, ResolutionFailed
+from hypothesis.errors import InvalidArgument
 from hypothesis.internal.compat import get_type_hints
 from hypothesis.internal.reflection import is_mock
 from hypothesis.internal.validation import check_type
-from hypothesis.strategies._internal.strategies import OneOfStrategy
+from hypothesis.strategies._internal.core import BuildsStrategy
+from hypothesis.strategies._internal.flatmapped import FlatMapStrategy
+from hypothesis.strategies._internal.lazy import LazyStrategy, unwrap_strategies
+from hypothesis.strategies._internal.strategies import (
+    FilteredStrategy,
+    MappedSearchStrategy,
+    OneOfStrategy,
+)
 from hypothesis.strategies._internal.types import _global_type_lookup
 from hypothesis.utils.conventions import InferType, infer
 
@@ -106,7 +113,7 @@
 # This test code was written by the `hypothesis.extra.ghostwriter` module
 # and is provided under the Creative Commons Zero public domain dedication.
 
-{imports}from hypothesis import given, {reject}strategies as st
+{imports}
 """
 
 TEMPLATE = """
@@ -271,11 +278,13 @@ def _get_strategies(
         with _with_any_registered():
             strat = st.builds(f, **builder_args).wrapped_strategy  # type: ignore
 
-        args, kwargs = strat.mapped_strategy.wrapped_strategy.element_strategies
-        if args.element_strategies:
+        if strat.args:
             raise NotImplementedError("Expected to pass everything as kwargs")
 
-        for k, v in zip(kwargs.keys, kwargs.mapped_strategy.element_strategies):
+        for k, v in strat.kwargs.items():
+            if _valid_syntax_repr(v)[1] == "nothing()" and k in hints:
+                # e.g. from_type(Hashable) is OK but the unwrapped repr is not
+                v = LazyStrategy(st.from_type, (hints[k],), {})
             if k in given_strategies:
                 given_strategies[k] |= v
             else:
@@ -297,10 +306,58 @@ def _assert_eq(style, a, b):
     return f"assert {a} == {b}"
 
 
+def _imports_for_object(obj):
+    """Return the imports for `obj`, which may be empty for e.g. lambdas"""
+    try:
+        if (not callable(obj)) or obj.__name__ == "<lambda>":
+            return set()
+        name = _get_qualname(obj).split(".")[0]
+        return {(_get_module(obj), name)}
+    except Exception:
+        return set()
+
+
+def _imports_for_strategy(strategy):
+    # If we have a lazy from_type strategy, because unwrapping it gives us an
+    # error or invalid syntax, import that type and we're done.
+    if isinstance(strategy, LazyStrategy) and strategy.function is st.from_type:
+        return _imports_for_object(strategy._LazyStrategy__args[0])
+
+    imports = set()
+    strategy = unwrap_strategies(strategy)
+
+    # Get imports for s.map(f), s.filter(f), s.flatmap(f), including both s and f
+    if isinstance(strategy, MappedSearchStrategy):
+        imports |= _imports_for_strategy(strategy.mapped_strategy)
+        imports |= _imports_for_object(strategy.pack)
+    if isinstance(strategy, FilteredStrategy):
+        imports |= _imports_for_strategy(strategy.filtered_strategy)
+        for f in strategy.flat_conditions:
+            imports |= _imports_for_object(f)
+    if isinstance(strategy, FlatMapStrategy):
+        imports |= _imports_for_strategy(strategy.flatmapped_strategy)
+        imports |= _imports_for_object(strategy.expand)
+
+    # recurse through one_of to handle e.g. from_type(Optional[Foo])
+    if isinstance(strategy, OneOfStrategy):
+        for s in strategy.element_strategies:
+            imports |= _imports_for_strategy(s)
+
+    # get imports for the target of builds(), and recurse into the argument strategies
+    if isinstance(strategy, BuildsStrategy):
+        imports |= _imports_for_object(strategy.target)
+        for s in strategy.args:
+            imports |= _imports_for_strategy(s)
+        for s in strategy.kwargs.values():
+            imports |= _imports_for_strategy(s)
+
+    return imports
+
+
 def _valid_syntax_repr(strategy):
     # For binary_op, we pass a variable name - so pass it right back again.
     if isinstance(strategy, str):
-        return strategy
+        return set(), strategy
     # Flatten and de-duplicate any one_of strategies, whether that's from resolving
     # a Union type or combining inputs to multiple functions.
     try:
@@ -314,17 +371,25 @@ def _valid_syntax_repr(strategy):
             strategy = st.one_of(elems or st.nothing())
         # Trivial special case because the wrapped repr for text() is terrible.
         if strategy == st.text().wrapped_strategy:
-            return "text()"
+            return set(), "text()"
         # Return a syntactically-valid strategy repr, including fixing some
         # strategy reprs and replacing invalid syntax reprs with `"nothing()"`.
         # String-replace to hide the special case in from_type() for Decimal('snan')
         r = repr(strategy).replace(".filter(_can_hash)", "")
+        # Replace <unknown> with ... in confusing lambdas
+        r = re.sub(r"(lambda.*?: )(<unknown>)([,)])", r"\1...\3", r)
         compile(r, "<string>", "eval")
-        return r
-    except (SyntaxError, ResolutionFailed):
-        return "nothing()"
+        # Finally, try to work out the imports we need for builds(), .map(),
+        # .filter(), and .flatmap() to work without NameError
+        imports = {i for i in _imports_for_strategy(strategy) if i[1] in r}
+        return imports, r
+    except (SyntaxError, InvalidArgument):
+        return set(), "nothing()"
 
 
+# When we ghostwrite for a module, we want to treat that as the __module__ for
+# each function, rather than whichever internal file it was actually defined in.
+KNOWN_FUNCTION_LOCATIONS: Dict[object, str] = {}
 # (g)ufuncs do not have a __module__ attribute, so we simply look for them in
 # any of the following modules which are found in sys.modules.  The order of
 # these entries doesn't matter, because we check identity of the found object.
@@ -332,6 +397,8 @@ def _valid_syntax_repr(strategy):
 
 
 def _get_module(obj):
+    if obj in KNOWN_FUNCTION_LOCATIONS:
+        return KNOWN_FUNCTION_LOCATIONS[obj]
     try:
         return obj.__module__
     except AttributeError:
@@ -374,23 +441,23 @@ def _make_test_body(
     except_: Tuple[Type[Exception], ...],
     style: str,
     given_strategies: Optional[Mapping[str, Union[str, st.SearchStrategy]]] = None,
-) -> Tuple[Set[str], str]:
+) -> Tuple[Set[Union[str, Tuple[str, str]]], str]:
+    # A set of modules to import - we might add to this later.  The import code
+    # is written later, so we can have one import section for multiple magic()
+    # test functions.
+    imports = {_get_module(f) for f in funcs}
+
     # Get strategies for all the arguments to each function we're testing.
     with _with_any_registered():
         given_strategies = given_strategies or _get_strategies(
             *funcs, pass_result_to_next_func=ghost in ("idempotent", "roundtrip")
         )
-        given_args = ", ".join(
-            f"{k}={_valid_syntax_repr(v)}" for k, v in given_strategies.items()
-        )
+        reprs = [((k,) + _valid_syntax_repr(v)) for k, v in given_strategies.items()]
+        imports = imports.union(*(imp for _, imp, _ in reprs))
+        given_args = ", ".join(f"{k}={v}" for k, _, v in reprs)
     for name in st.__all__:
         given_args = given_args.replace(f"{name}(", f"st.{name}(")
 
-    # A set of modules to import - we might add to this later.  The import code
-    # is written later, so we can have one import section for multiple magic()
-    # test functions.
-    imports = {_get_module(f) for f in funcs}
-
     if except_:
         # This is reminiscent of de-duplication logic I wrote for flake8-bugbear,
         # but with access to the actual objects we can just check for subclasses.
@@ -438,18 +505,27 @@ def _make_test_body(
     return imports, body
 
 
-def _make_test(imports: Set[str], body: str) -> str:
+def _make_test(imports: Set[Union[str, Tuple[str, str]]], body: str) -> str:
     # Discarding "builtins." and "__main__" probably isn't particularly useful
     # for user code, but important for making a good impression in demos.
     body = body.replace("builtins.", "").replace("__main__.", "")
-    imports.discard("builtins")
-    imports.discard("__main__")
     if "st.from_type(typing." in body:
         imports.add("typing")
-    header = IMPORT_SECTION.format(
-        imports="".join(f"import {imp}\n" for imp in sorted(imports)),
-        reject="reject, " if "        reject()\n" in body else "",
-    )
+    imports |= {("hypothesis", "given"), ("hypothesis", "strategies as st")}
+    if "        reject()\n" in body:
+        imports.add(("hypothesis", "reject"))
+
+    do_not_import = {"builtins", "__main__"}
+    direct = {f"import {i}" for i in imports - do_not_import if isinstance(i, str)}
+    from_imports = defaultdict(set)
+    for module, name in {i for i in imports if isinstance(i, tuple)}:
+        from_imports[module].add(name)
+    from_ = {
+        "from {} import {}".format(module, ", ".join(sorted(names)))
+        for module, names in from_imports.items()
+        if module not in do_not_import
+    }
+    header = IMPORT_SECTION.format(imports="\n".join(sorted(direct) + sorted(from_)))
     nothings = body.count("st.nothing()")
     if nothings == 1:
         header += "# TODO: replace st.nothing() with an appropriate strategy\n\n"
@@ -525,8 +601,15 @@ def magic(
                 ]
             for f in funcs:
                 try:
-                    if (not is_mock(f)) and callable(f) and _get_params(f):
+                    if (
+                        (not is_mock(f))
+                        and callable(f)
+                        and _get_params(f)
+                        and not isinstance(f, enum.EnumMeta)
+                    ):
                         functions.add(f)
+                        if getattr(thing, "__name__", None):
+                            KNOWN_FUNCTION_LOCATIONS[f] = thing.__name__
                 except (TypeError, ValueError):
                     pass
         else:
@@ -866,7 +949,7 @@ def _make_binop_body(
     distributes_over: Optional[Callable[[X, X], X]] = None,
     except_: Tuple[Type[Exception], ...],
     style: str,
-) -> Tuple[Set[str], str]:
+) -> Tuple[Set[Union[str, Tuple[str, str]]], str]:
     # TODO: collapse togther first two strategies, keep any others (for flags etc.)
     # assign this as a global variable, which will be prepended to the test bodies
     strategies = _get_strategies(func)
@@ -951,7 +1034,7 @@ def maker(
             _write_call(func, "a", _write_call(distributes_over, "b", "c")),
         )
 
-    operands_repr = _valid_syntax_repr(operands)
+    _, operands_repr = _valid_syntax_repr(operands)
     for name in st.__all__:
         operands_repr = operands_repr.replace(f"{name}(", f"st.{name}(")
     classdef = ""
diff --git a/hypothesis-python/src/hypothesis/strategies/_internal/core.py b/hypothesis-python/src/hypothesis/strategies/_internal/core.py
--- a/hypothesis-python/src/hypothesis/strategies/_internal/core.py
+++ b/hypothesis-python/src/hypothesis/strategies/_internal/core.py
@@ -1282,6 +1282,36 @@ def random_module() -> SearchStrategy[RandomSeeder]:
     return shared(RandomModule(), key="hypothesis.strategies.random_module()")
 
 
+class BuildsStrategy(SearchStrategy):
+    def __init__(self, target, args, kwargs):
+        self.target = target
+        self.args = args
+        self.kwargs = kwargs
+
+    def do_draw(self, data):
+        try:
+            return self.target(
+                *[data.draw(a) for a in self.args],
+                **{k: data.draw(v) for k, v in self.kwargs.items()},
+            )
+        except TypeError as err:
+            if (
+                isinstance(self.target, type)
+                and issubclass(self.target, enum.Enum)
+                and not (self.args or self.kwargs)
+            ):
+                name = self.target.__module__ + "." + self.target.__qualname__
+                raise InvalidArgument(
+                    f"Calling {name} with no arguments raised an error - "
+                    f"try using sampled_from({name}) instead of builds({name})"
+                ) from err
+            raise
+
+    def validate(self):
+        tuples(*self.args).validate()
+        fixed_dictionaries(self.kwargs).validate()
+
+
 # The ideal signature builds(target, /, *args, **kwargs) is unfortunately a
 # SyntaxError before Python 3.8 so we emulate it with manual argument unpacking.
 # Note that for the benefit of documentation and introspection tools, we set the
@@ -1349,27 +1379,9 @@ def builds(
         for kw in set(hints) & (required | to_infer):
             kwargs[kw] = from_type(hints[kw])
 
-    def build_target(value):
-        args, kwargs = value
-        try:
-            return target(*args, **kwargs)
-        except TypeError as err:
-            if (
-                isinstance(target, type)
-                and issubclass(target, enum.Enum)
-                and not (args or kwargs)
-            ):
-                name = target.__module__ + "." + target.__qualname__
-                raise InvalidArgument(
-                    f"Calling {name} with no arguments raised an error - "
-                    f"try using sampled_from({name}) instead of builds({name})"
-                ) from err
-            raise
-
     # Mypy doesn't realise that `infer` is gone from kwargs now
     # and thinks that target and args have the same (union) type.
-    args_kwargs = tuples(tuples(*args), fixed_dictionaries(kwargs))  # type: ignore
-    return args_kwargs.map(build_target)
+    return BuildsStrategy(target, args, kwargs)
 
 
 if sys.version_info[:2] >= (3, 8):  # pragma: no cover


Test patch:
diff --git a/hypothesis-python/tests/ghostwriter/recorded/hypothesis_module_magic.txt b/hypothesis-python/tests/ghostwriter/recorded/hypothesis_module_magic.txt
new file mode 100644
--- /dev/null
+++ b/hypothesis-python/tests/ghostwriter/recorded/hypothesis_module_magic.txt
@@ -0,0 +1,100 @@
+# This test code was written by the `hypothesis.extra.ghostwriter` module
+# and is provided under the Creative Commons Zero public domain dedication.
+
+import hypothesis
+import typing
+from hypothesis import given, settings, strategies as st
+from random import Random
+
+
+@given(condition=st.from_type(object))
+def test_fuzz_assume(condition):
+    hypothesis.assume(condition=condition)
+
+
+@given(value=st.text())
+def test_fuzz_event(value):
+    hypothesis.event(value=value)
+
+
+@given(
+    specifier=st.from_type(hypothesis.strategies.SearchStrategy),
+    condition=st.functions(like=lambda *a: ..., returns=st.booleans()),
+    settings=st.one_of(st.none(), st.builds(settings)),
+    random=st.one_of(st.none(), st.builds(Random)),
+    database_key=st.one_of(st.none(), st.binary()),
+)
+def test_fuzz_find(specifier, condition, settings, random, database_key):
+    hypothesis.find(
+        specifier=specifier,
+        condition=condition,
+        settings=settings,
+        random=random,
+        database_key=database_key,
+    )
+
+
+@given(value=st.text())
+def test_fuzz_note(value):
+    hypothesis.note(value=value)
+
+
+@given(r=st.builds(Random))
+def test_fuzz_register_random(r):
+    hypothesis.register_random(r=r)
+
+
+@given(version=st.text(), blob=st.binary())
+def test_fuzz_reproduce_failure(version, blob):
+    hypothesis.reproduce_failure(version=version, blob=blob)
+
+
+@given(seed=st.from_type(typing.Hashable))
+def test_fuzz_seed(seed):
+    hypothesis.seed(seed=seed)
+
+
+@given(
+    parent=st.none(),
+    max_examples=st.just(not_set),
+    derandomize=st.just(not_set),
+    database=st.just(not_set),
+    verbosity=st.just(not_set),
+    phases=st.just(not_set),
+    stateful_step_count=st.just(not_set),
+    report_multiple_bugs=st.just(not_set),
+    suppress_health_check=st.just(not_set),
+    deadline=st.just(not_set),
+    print_blob=st.just(not_set),
+)
+def test_fuzz_settings(
+    parent,
+    max_examples,
+    derandomize,
+    database,
+    verbosity,
+    phases,
+    stateful_step_count,
+    report_multiple_bugs,
+    suppress_health_check,
+    deadline,
+    print_blob,
+):
+    hypothesis.settings(
+        parent=parent,
+        max_examples=max_examples,
+        derandomize=derandomize,
+        database=database,
+        verbosity=verbosity,
+        phases=phases,
+        stateful_step_count=stateful_step_count,
+        report_multiple_bugs=report_multiple_bugs,
+        suppress_health_check=suppress_health_check,
+        deadline=deadline,
+        print_blob=print_blob,
+    )
+
+
+@given(observation=st.one_of(st.floats(), st.integers()), label=st.text())
+def test_fuzz_target(observation, label):
+    hypothesis.target(observation=observation, label=label)
diff --git a/hypothesis-python/tests/ghostwriter/test_expected_output.py b/hypothesis-python/tests/ghostwriter/test_expected_output.py
--- a/hypothesis-python/tests/ghostwriter/test_expected_output.py
+++ b/hypothesis-python/tests/ghostwriter/test_expected_output.py
@@ -24,12 +24,15 @@
 import operator
 import pathlib
 import re
+import sys
 from typing import Sequence
 
 import numpy
 import pytest
 
+import hypothesis
 from hypothesis.extra import ghostwriter
+from hypothesis.utils.conventions import not_set
 
 
 @pytest.fixture
@@ -37,6 +40,13 @@ def update_recorded_outputs(request):
     return request.config.getoption("--hypothesis-update-outputs")
 
 
+def get_recorded(name, actual=""):
+    file_ = pathlib.Path(__file__).parent / "recorded" / f"{name}.txt"
+    if actual:
+        file_.write_text(actual)
+    return file_.read_text()
+
+
 def timsort(seq: Sequence[int]) -> Sequence[int]:
     return sorted(seq)
 
@@ -101,9 +111,16 @@ def add(a: float, b: float) -> float:
 )
 def test_ghostwriter_example_outputs(update_recorded_outputs, data):
     name, actual = data
-    file_ = pathlib.Path(__file__).parent / "recorded" / f"{name}.txt"
-    if update_recorded_outputs:
-        file_.write_text(actual)
-    expected = file_.read_text()
+    expected = get_recorded(name, actual * update_recorded_outputs)
     assert actual == expected  # We got the expected source code
     exec(expected, {})  # and there are no SyntaxError or NameErrors
+
+
+def test_ghostwriter_on_hypothesis(update_recorded_outputs):
+    actual = ghostwriter.magic(hypothesis).replace("Strategy[+Ex]", "Strategy")
+    expected = get_recorded("hypothesis_module_magic", actual * update_recorded_outputs)
+    # The py36 typing module has some different handling of generics (SearchStrategy)
+    # and contents (collections.abc vs typing), but we can still check the code works.
+    if sys.version_info[:2] > (3, 6):
+        assert actual == expected
+    exec(expected, {"not_set": not_set})
diff --git a/hypothesis-python/tests/ghostwriter/test_ghostwriter.py b/hypothesis-python/tests/ghostwriter/test_ghostwriter.py
--- a/hypothesis-python/tests/ghostwriter/test_ghostwriter.py
+++ b/hypothesis-python/tests/ghostwriter/test_ghostwriter.py
@@ -27,7 +27,8 @@
 
 from hypothesis.errors import InvalidArgument, MultipleFailures, Unsatisfiable
 from hypothesis.extra import ghostwriter
-from hypothesis.strategies import from_type, just
+from hypothesis.strategies import builds, from_type, just
+from hypothesis.strategies._internal.lazy import LazyStrategy
 
 varied_excepts = pytest.mark.parametrize("ex", [(), ValueError, (TypeError, re.error)])
 
@@ -71,7 +72,7 @@ def __repr__(self):
 
     s = just(NoRepr())
     assert repr(s) == f"just({msg})"
-    assert ghostwriter._valid_syntax_repr(s) == "nothing()"
+    assert ghostwriter._valid_syntax_repr(s)[1] == "nothing()"
 
 
 class AnEnum(enum.Enum):
@@ -119,7 +120,7 @@ def non_resolvable_arg(x: NotResolvable):
 def test_flattens_one_of_repr():
     strat = from_type(Union[int, Sequence[int]])
     assert repr(strat).count("one_of(") > 1
-    assert ghostwriter._valid_syntax_repr(strat).count("one_of(") == 1
+    assert ghostwriter._valid_syntax_repr(strat)[1].count("one_of(") == 1
 
 
 @varied_excepts
@@ -290,3 +291,39 @@ def test_unrepr_identity_elem():
     # and also works with explicit identity element
     source_code = ghostwriter.binary_operation(compose_types, identity=type)
     exec(source_code, {})
+
+
+@pytest.mark.parametrize(
+    "strategy, imports",
+    # The specifics don't matter much here; we're just demonstrating that
+    # we can walk the strategy and collect all the objects to import.
+    [
+        # Lazy from_type() is handled without being unwrapped
+        (LazyStrategy(from_type, (enum.Enum,), {}), {("enum", "Enum")}),
+        # Mapped, filtered, and flatmapped check both sides of the method
+        (
+            builds(enum.Enum).map(Decimal),
+            {("enum", "Enum"), ("decimal", "Decimal")},
+        ),
+        (
+            builds(enum.Enum).flatmap(Decimal),
+            {("enum", "Enum"), ("decimal", "Decimal")},
+        ),
+        (
+            builds(enum.Enum).filter(Decimal).filter(re.compile),
+            {("enum", "Enum"), ("decimal", "Decimal"), ("re", "compile")},
+        ),
+        # one_of() strategies recurse into all the branches
+        (
+            builds(enum.Enum) | builds(Decimal) | builds(re.compile),
+            {("enum", "Enum"), ("decimal", "Decimal"), ("re", "compile")},
+        ),
+        # and builds() checks the arguments as well as the target
+        (
+            builds(enum.Enum, builds(Decimal), kw=builds(re.compile)),
+            {("enum", "Enum"), ("decimal", "Decimal"), ("re", "compile")},
+        ),
+    ],
+)
+def test_get_imports_for_strategy(strategy, imports):
+    assert ghostwriter._imports_for_strategy(strategy) == imports