Repository: HypothesisWorks/hypothesis

Issue:
Add type hints to the tests generated by ghostwriter
I was wondering if it is possible to add type hints to the tests generated by ghostwriter. For that I've come up with a straight-forward solution (~44 insertions), which uses the `_get_params` method to get the parameters and uses them to see what annotations are required. If `inspect.Parameter.empty`  is encountered the parameter does not get any type hints.

The solution is probably far from perfect, but it might make creating the tests easier if type hints are required in a project.

Are you interested in that? If so I can create a PR for it. If it is too much change in behaviour I could also add an option to the `write` CLI command with which this could be enabled.


Discussion:
Seems like a reasonable feature!  I think we'll want to add add an `annotations: bool | None = None` argument: if boolean we do or don't write out type annotations, and if None we match the annotated-ness of the code we're writing tests for.  A PR would be fantastic; I'm very happy to review things and help contributors get unstuck.

Resolving patch:
diff --git a/hypothesis-python/src/hypothesis/extra/cli.py b/hypothesis-python/src/hypothesis/extra/cli.py
--- a/hypothesis-python/src/hypothesis/extra/cli.py
+++ b/hypothesis-python/src/hypothesis/extra/cli.py
@@ -269,7 +269,15 @@ def codemod(path):
         multiple=True,
         help="dotted name of exception(s) to ignore",
     )
-    def write(func, writer, except_, style):  # noqa: D301  # \b disables autowrap
+    @click.option(
+        "--annotate/--no-annotate",
+        default=None,
+        help="force ghostwritten tests to be type-annotated (or not).  "
+        "By default, match the code to test.",
+    )
+    def write(
+        func, writer, except_, style, annotate
+    ):  # noqa: D301  # \b disables autowrap
         """`hypothesis write` writes property-based tests for you!
 
         Type annotations are helpful but not required for our advanced introspection
@@ -278,6 +286,7 @@ def write(func, writer, except_, style):  # noqa: D301  # \b disables autowrap
         \b
             hypothesis write gzip
             hypothesis write numpy.matmul
+            hypothesis write pandas.from_dummies
             hypothesis write re.compile --except re.error
             hypothesis write --equivalent ast.literal_eval eval
             hypothesis write --roundtrip json.dumps json.loads
@@ -287,7 +296,7 @@ def write(func, writer, except_, style):  # noqa: D301  # \b disables autowrap
         # NOTE: if you want to call this function from Python, look instead at the
         # ``hypothesis.extra.ghostwriter`` module.  Click-decorated functions have
         # a different calling convention, and raise SystemExit instead of returning.
-        kwargs = {"except_": except_ or (), "style": style}
+        kwargs = {"except_": except_ or (), "style": style, "annotate": annotate}
         if writer is None:
             writer = "magic"
         elif writer == "idempotent" and len(func) > 1:
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
@@ -33,6 +33,7 @@
 
           hypothesis write gzip
           hypothesis write numpy.matmul
+          hypothesis write pandas.from_dummies
           hypothesis write re.compile --except re.error
           hypothesis write --equivalent ast.literal_eval eval
           hypothesis write --roundtrip json.dumps json.loads
@@ -40,14 +41,16 @@
           hypothesis write --binary-op operator.add
 
     Options:
-      --roundtrip                start by testing write/read or encode/decode!
-      --equivalent               very useful when optimising or refactoring code
-      --errors-equivalent        --equivalent, but also allows consistent errors
-      --idempotent               check that f(x) == f(f(x))
-      --binary-op                associativity, commutativity, identity element
-      --style [pytest|unittest]  pytest-style function, or unittest-style method?
-      -e, --except OBJ_NAME      dotted name of exception(s) to ignore
-      -h, --help                 Show this message and exit.
+      --roundtrip                 start by testing write/read or encode/decode!
+      --equivalent                very useful when optimising or refactoring code
+      --errors-equivalent         --equivalent, but also allows consistent errors
+      --idempotent                check that f(x) == f(f(x))
+      --binary-op                 associativity, commutativity, identity element
+      --style [pytest|unittest]   pytest-style function, or unittest-style method?
+      -e, --except OBJ_NAME       dotted name of exception(s) to ignore
+      --annotate / --no-annotate  force ghostwritten tests to be type-annotated
+                                  (or not).  By default, match the code to test.
+      -h, --help                  Show this message and exit.
 
 .. tip::
 
@@ -86,9 +89,13 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    DefaultDict,
     Dict,
+    ForwardRef,
+    Iterable,
     List,
     Mapping,
+    NamedTuple,
     Optional,
     Set,
     Tuple,
@@ -101,7 +108,7 @@
 
 from hypothesis import Verbosity, find, settings, strategies as st
 from hypothesis.errors import InvalidArgument
-from hypothesis.internal.compat import get_type_hints
+from hypothesis.internal.compat import get_args, get_origin, get_type_hints
 from hypothesis.internal.reflection import get_signature, is_mock
 from hypothesis.internal.validation import check_type
 from hypothesis.provisional import domains
@@ -134,7 +141,7 @@
 
 TEMPLATE = """
 @given({given_args})
-def test_{test_kind}_{func_name}({arg_names}):
+def test_{test_kind}_{func_name}({arg_names}){return_annotation}:
 {test_body}
 """
 
@@ -760,6 +767,7 @@ def _make_test_body(
     style: str,
     given_strategies: Optional[Mapping[str, Union[str, st.SearchStrategy]]] = None,
     imports: Optional[ImportSet] = None,
+    annotate: bool,
 ) -> Tuple[ImportSet, 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()
@@ -791,12 +799,18 @@ def _make_test_body(
         test_body = f"{test_body}\n{assertions}"
 
     # Indent our test code to form the body of a function or method.
-    argnames = (["self"] if style == "unittest" else []) + list(given_strategies)
+    argnames = ["self"] if style == "unittest" else []
+    if annotate:
+        argnames.extend(_annotate_args(given_strategies, funcs, imports))
+    else:
+        argnames.extend(given_strategies)
+
     body = TEMPLATE.format(
         given_args=given_args,
         test_kind=ghost,
         func_name="_".join(_get_qualname(f).replace(".", "_") for f in funcs),
         arg_names=", ".join(argnames),
+        return_annotation=" -> None" if annotate else "",
         test_body=indent(test_body, prefix="    "),
     )
 
@@ -812,6 +826,160 @@ def _make_test_body(
     return imports, body
 
 
+def _annotate_args(
+    argnames: Iterable[str], funcs: Iterable[Callable], imports: ImportSet
+) -> Iterable[str]:
+    arg_parameters: DefaultDict[str, Set[Any]] = defaultdict(set)
+    for func in funcs:
+        for key, param in _get_params(func).items():
+            arg_parameters[key].add(param.annotation)
+
+    for argname in argnames:
+        parameters = arg_parameters.get(argname)
+        annotation = (
+            None
+            if parameters is None
+            else _parameters_to_annotation_name(parameters, imports)
+        )
+        if annotation is None:
+            yield argname
+        else:
+            yield f"{argname}: {annotation}"
+
+
+class _AnnotationData(NamedTuple):
+    type_name: str
+    imports: Set[str]
+
+
+def _parameters_to_annotation_name(
+    parameters: Iterable[Any], imports: ImportSet
+) -> Optional[str]:
+    annotations = tuple(
+        annotation
+        for annotation in (
+            _parameter_to_annotation(parameter)
+            for parameter in parameters
+            if parameter != inspect.Parameter.empty
+        )
+        if annotation is not None
+    )
+    if not annotations:
+        return None
+    if len(annotations) == 1:
+        type_name, new_imports = annotations[0]
+        imports.update(new_imports)
+        return type_name
+    joined = _join_generics(("typing.Union", {"typing"}), annotations)
+    if joined is None:
+        return None
+    imports.update(joined.imports)
+    return joined.type_name
+
+
+def _join_generics(
+    origin_type_data: Optional[Tuple[str, Set[str]]],
+    annotations: Iterable[Optional[_AnnotationData]],
+) -> Optional[_AnnotationData]:
+    if origin_type_data is None:
+        return None
+
+    origin_type, imports = origin_type_data
+    joined = _join_argument_annotations(annotations)
+    if joined is None or not joined[0]:
+        return None
+
+    arg_types, new_imports = joined
+    imports.update(new_imports)
+    return _AnnotationData("{}[{}]".format(origin_type, ", ".join(arg_types)), imports)
+
+
+def _join_argument_annotations(
+    annotations: Iterable[Optional[_AnnotationData]],
+) -> Optional[Tuple[List[str], Set[str]]]:
+    imports: Set[str] = set()
+    arg_types: List[str] = []
+
+    for annotation in annotations:
+        if annotation is None:
+            return None
+        arg_types.append(annotation.type_name)
+        imports.update(annotation.imports)
+
+    return arg_types, imports
+
+
+def _parameter_to_annotation(parameter: Any) -> Optional[_AnnotationData]:
+    # if a ForwardRef could not be resolved
+    if isinstance(parameter, str):
+        return None
+
+    if isinstance(parameter, ForwardRef):
+        forwarded_value = parameter.__forward_value__
+        if forwarded_value is None:
+            return None
+        return _parameter_to_annotation(forwarded_value)
+
+    if isinstance(parameter, type):
+        if parameter.__module__ == "builtins":
+            return _AnnotationData(
+                "None" if parameter.__name__ == "NoneType" else parameter.__name__,
+                set(),
+            )
+        return _AnnotationData(
+            f"{parameter.__module__}.{parameter.__name__}", {parameter.__module__}
+        )
+
+    # the arguments of Callable are in a list
+    if isinstance(parameter, list):
+        joined = _join_argument_annotations(map(_parameter_to_annotation, parameter))
+        if joined is None:
+            return None
+        arg_type_names, new_imports = joined
+        return _AnnotationData("[{}]".format(", ".join(arg_type_names)), new_imports)
+
+    origin_type = get_origin(parameter)
+
+    # if not generic or no generic arguments
+    if origin_type is None or origin_type == parameter:
+        type_name = str(parameter)
+        if type_name.startswith("typing."):
+            return _AnnotationData(type_name, {"typing"})
+        return _AnnotationData(type_name, set())
+
+    arg_types = get_args(parameter)
+    type_name = str(parameter)
+
+    # typing types get translated to classes that don't support generics
+    origin_annotation: Optional[_AnnotationData]
+    if type_name.startswith("typing."):
+        try:
+            new_type_name = type_name[: type_name.index("[")]
+        except ValueError:
+            new_type_name = type_name
+        origin_annotation = _AnnotationData(new_type_name, {"typing"})
+    else:
+        origin_annotation = _parameter_to_annotation(origin_type)
+
+    if arg_types:
+        return _join_generics(
+            origin_annotation, map(_parameter_to_annotation, arg_types)
+        )
+    return origin_annotation
+
+
+def _are_annotations_used(*functions: Callable) -> bool:
+    for function in functions:
+        try:
+            params = get_signature(function).parameters.values()
+        except Exception:
+            pass
+        else:
+            if any(param.annotation != inspect.Parameter.empty for param in params):
+                return True
+    return False
+
+
 def _make_test(imports: ImportSet, body: str) -> str:
     # Discarding "builtins." and "__main__" probably isn't particularly useful
     # for user code, but important for making a good impression in demos.
@@ -880,6 +1048,7 @@ def magic(
     *modules_or_functions: Union[Callable, types.ModuleType],
     except_: Except = (),
     style: str = "pytest",
+    annotate: Optional[bool] = None,
 ) -> str:
     """Guess which ghostwriters to use, for a module or collection of functions.
 
@@ -952,6 +1121,9 @@ def magic(
             except (TypeError, ValueError):
                 pass
 
+    if annotate is None:
+        annotate = _are_annotations_used(*functions)
+
     imports = set()
     parts = []
 
@@ -984,7 +1156,11 @@ def make_(how, *args, **kwargs):
                 for other in sorted(
                     n for n in by_name if n.split(".")[-1] == inverse_name
                 ):
-                    make_(_make_roundtrip_body, (by_name.pop(name), by_name.pop(other)))
+                    make_(
+                        _make_roundtrip_body,
+                        (by_name.pop(name), by_name.pop(other)),
+                        annotate=annotate,
+                    )
                     break
                 else:
                     try:
@@ -996,7 +1172,11 @@ def make_(how, *args, **kwargs):
                     except Exception:
                         pass
                     else:
-                        make_(_make_roundtrip_body, (by_name.pop(name), other_func))
+                        make_(
+                            _make_roundtrip_body,
+                            (by_name.pop(name), other_func),
+                            annotate=annotate,
+                        )
 
     # Look for equivalent functions: same name, all required arguments of any can
     # be found in all signatures, and if all have return-type annotations they match.
@@ -1008,7 +1188,7 @@ def make_(how, *args, **kwargs):
             sentinel = object()
             returns = {get_type_hints(f).get("return", sentinel) for f in group}
             if len(returns - {sentinel}) <= 1:
-                make_(_make_equiv_body, group)
+                make_(_make_equiv_body, group, annotate=annotate)
                 for f in group:
                     by_name.pop(_get_qualname(f, include_module=True))
 
@@ -1022,14 +1202,14 @@ def make_(how, *args, **kwargs):
             a, b = hints.values()
             arg1, arg2 = params
             if a == b and len(arg1) == len(arg2) <= 3:
-                make_(_make_binop_body, func)
+                make_(_make_binop_body, func, annotate=annotate)
                 del by_name[name]
 
     # Look for Numpy ufuncs or gufuncs, and write array-oriented tests for them.
     if "numpy" in sys.modules:
         for name, func in sorted(by_name.items()):
             if _is_probably_ufunc(func):
-                make_(_make_ufunc_body, func)
+                make_(_make_ufunc_body, func, annotate=annotate)
                 del by_name[name]
 
     # For all remaining callables, just write a fuzz-test.  In principle we could
@@ -1037,12 +1217,22 @@ def make_(how, *args, **kwargs):
     # be worth the trouble when it's so easy for the user to specify themselves.
     for _, f in sorted(by_name.items()):
         make_(
-            _make_test_body, f, test_body=_write_call(f, except_=except_), ghost="fuzz"
+            _make_test_body,
+            f,
+            test_body=_write_call(f, except_=except_),
+            ghost="fuzz",
+            annotate=annotate,
         )
     return _make_test(imports, "\n".join(parts))
 
 
-def fuzz(func: Callable, *, except_: Except = (), style: str = "pytest") -> str:
+def fuzz(
+    func: Callable,
+    *,
+    except_: Except = (),
+    style: str = "pytest",
+    annotate: Optional[bool] = None,
+) -> str:
     """Write source code for a property-based test of ``func``.
 
     The resulting test checks that valid input only leads to expected exceptions.
@@ -1086,17 +1276,27 @@ def test_fuzz_compile(pattern, flags):
     except_ = _check_except(except_)
     _check_style(style)
 
+    if annotate is None:
+        annotate = _are_annotations_used(func)
+
     imports, body = _make_test_body(
         func,
         test_body=_write_call(func, except_=except_),
         except_=except_,
         ghost="fuzz",
         style=style,
+        annotate=annotate,
     )
     return _make_test(imports, body)
 
 
-def idempotent(func: Callable, *, except_: Except = (), style: str = "pytest") -> str:
+def idempotent(
+    func: Callable,
+    *,
+    except_: Except = (),
+    style: str = "pytest",
+    annotate: Optional[bool] = None,
+) -> str:
     """Write source code for a property-based test of ``func``.
 
     The resulting test checks that if you call ``func`` on it's own output,
@@ -1136,6 +1336,9 @@ def test_idempotent_timsort(seq):
     except_ = _check_except(except_)
     _check_style(style)
 
+    if annotate is None:
+        annotate = _are_annotations_used(func)
+
     imports, body = _make_test_body(
         func,
         test_body="result = {}\nrepeat = {}".format(
@@ -1146,11 +1349,12 @@ def test_idempotent_timsort(seq):
         assertions=_assert_eq(style, "result", "repeat"),
         ghost="idempotent",
         style=style,
+        annotate=annotate,
     )
     return _make_test(imports, body)
 
 
-def _make_roundtrip_body(funcs, except_, style):
+def _make_roundtrip_body(funcs, except_, style, annotate):
     first_param = next(iter(_get_params(funcs[0])))
     test_lines = [
         _write_call(funcs[0], assign="value0", except_=except_),
@@ -1166,10 +1370,16 @@ def _make_roundtrip_body(funcs, except_, style):
         assertions=_assert_eq(style, first_param, f"value{len(funcs) - 1}"),
         ghost="roundtrip",
         style=style,
+        annotate=annotate,
     )
 
 
-def roundtrip(*funcs: Callable, except_: Except = (), style: str = "pytest") -> str:
+def roundtrip(
+    *funcs: Callable,
+    except_: Except = (),
+    style: str = "pytest",
+    annotate: Optional[bool] = None,
+) -> str:
     """Write source code for a property-based test of ``funcs``.
 
     The resulting test checks that if you call the first function, pass the result
@@ -1190,10 +1400,14 @@ def roundtrip(*funcs: Callable, except_: Except = (), style: str = "pytest") ->
             raise InvalidArgument(f"Got non-callable funcs[{i}]={f!r}")
     except_ = _check_except(except_)
     _check_style(style)
-    return _make_test(*_make_roundtrip_body(funcs, except_, style))
 
+    if annotate is None:
+        annotate = _are_annotations_used(*funcs)
 
-def _make_equiv_body(funcs, except_, style):
+    return _make_test(*_make_roundtrip_body(funcs, except_, style, annotate))
+
+
+def _make_equiv_body(funcs, except_, style, annotate):
     var_names = [f"result_{f.__name__}" for f in funcs]
     if len(set(var_names)) < len(var_names):
         var_names = [f"result_{i}_{ f.__name__}" for i, f in enumerate(funcs)]
@@ -1212,6 +1426,7 @@ def _make_equiv_body(funcs, except_, style):
         assertions=assertions,
         ghost="equivalent",
         style=style,
+        annotate=annotate,
     )
 
 
@@ -1234,7 +1449,7 @@ def _make_equiv_body(funcs, except_, style):
 """.rstrip()
 
 
-def _make_equiv_errors_body(funcs, except_, style):
+def _make_equiv_errors_body(funcs, except_, style, annotate):
     var_names = [f"result_{f.__name__}" for f in funcs]
     if len(set(var_names)) < len(var_names):
         var_names = [f"result_{i}_{ f.__name__}" for i, f in enumerate(funcs)]
@@ -1267,6 +1482,7 @@ def _make_equiv_errors_body(funcs, except_, style):
         except_=(),
         ghost="equivalent",
         style=style,
+        annotate=annotate,
     )
     return imports | extra_imports, source_code
 
@@ -1276,6 +1492,7 @@ def equivalent(
     allow_same_errors: bool = False,
     except_: Except = (),
     style: str = "pytest",
+    annotate: Optional[bool] = None,
 ) -> str:
     """Write source code for a property-based test of ``funcs``.
 
@@ -1303,10 +1520,14 @@ def equivalent(
     check_type(bool, allow_same_errors, "allow_same_errors")
     except_ = _check_except(except_)
     _check_style(style)
+
+    if annotate is None:
+        annotate = _are_annotations_used(*funcs)
+
     if allow_same_errors and not any(issubclass(Exception, ex) for ex in except_):
-        imports, source_code = _make_equiv_errors_body(funcs, except_, style)
+        imports, source_code = _make_equiv_errors_body(funcs, except_, style, annotate)
     else:
-        imports, source_code = _make_equiv_body(funcs, except_, style)
+        imports, source_code = _make_equiv_body(funcs, except_, style, annotate)
     return _make_test(imports, source_code)
 
 
@@ -1323,6 +1544,7 @@ def binary_operation(
     distributes_over: Optional[Callable[[X, X], X]] = None,
     except_: Except = (),
     style: str = "pytest",
+    annotate: Optional[bool] = None,
 ) -> str:
     """Write property tests for the binary operation ``func``.
 
@@ -1363,6 +1585,10 @@ def binary_operation(
         raise InvalidArgument(
             "You must select at least one property of the binary operation to test."
         )
+
+    if annotate is None:
+        annotate = _are_annotations_used(func)
+
     imports, body = _make_binop_body(
         func,
         associative=associative,
@@ -1371,6 +1597,7 @@ def binary_operation(
         distributes_over=distributes_over,
         except_=except_,
         style=style,
+        annotate=annotate,
     )
     return _make_test(imports, body)
 
@@ -1384,6 +1611,7 @@ def _make_binop_body(
     distributes_over: Optional[Callable[[X, X], X]] = None,
     except_: Tuple[Type[Exception], ...],
     style: str,
+    annotate: bool,
 ) -> Tuple[ImportSet, str]:
     strategies = _get_strategies(func)
     operands, b = (strategies.pop(p) for p in list(_get_params(func))[:2])
@@ -1413,6 +1641,7 @@ def maker(
             assertions=assertions,
             style=style,
             given_strategies={**strategies, **{n: operands_name for n in args}},
+            annotate=annotate,
         )
         all_imports.update(imports)
         if style == "unittest":
@@ -1506,7 +1735,13 @@ def maker(
     )
 
 
-def ufunc(func: Callable, *, except_: Except = (), style: str = "pytest") -> str:
+def ufunc(
+    func: Callable,
+    *,
+    except_: Except = (),
+    style: str = "pytest",
+    annotate: Optional[bool] = None,
+) -> str:
     """Write a property-based test for the :doc:`array ufunc <numpy:reference/ufuncs>` ``func``.
 
     The resulting test checks that your ufunc or :doc:`gufunc
@@ -1522,10 +1757,16 @@ def ufunc(func: Callable, *, except_: Except = (), style: str = "pytest") -> str
         raise InvalidArgument(f"func={func!r} does not seem to be a ufunc")
     except_ = _check_except(except_)
     _check_style(style)
-    return _make_test(*_make_ufunc_body(func, except_=except_, style=style))
+
+    if annotate is None:
+        annotate = _are_annotations_used(func)
+
+    return _make_test(
+        *_make_ufunc_body(func, except_=except_, style=style, annotate=annotate)
+    )
 
 
-def _make_ufunc_body(func, *, except_, style):
+def _make_ufunc_body(func, *, except_, style, annotate):
 
     import hypothesis.extra.numpy as npst
 
@@ -1570,4 +1811,5 @@ def _make_ufunc_body(func, *, except_, style):
         style=style,
         given_strategies={"data": st.data(), "shapes": shapes, "types": types},
         imports={("hypothesis.extra.numpy", "arrays")},
+        annotate=annotate,
     )
diff --git a/hypothesis-python/src/hypothesis/internal/compat.py b/hypothesis-python/src/hypothesis/internal/compat.py
--- a/hypothesis-python/src/hypothesis/internal/compat.py
+++ b/hypothesis-python/src/hypothesis/internal/compat.py
@@ -17,7 +17,7 @@
 from typing import Any, ForwardRef, Tuple
 
 try:
-    from typing import get_args
+    from typing import get_args as get_args
 except ImportError:
     # remove at Python 3.7 end-of-life
     from collections.abc import Callable as _Callable
@@ -46,6 +46,31 @@ def get_args(
         return ()
 
 
+try:
+    from typing import get_origin as get_origin
+except ImportError:
+    # remove at Python 3.7 end-of-life
+    from collections.abc import Callable as _Callable
+
+    def get_origin(tp: Any) -> typing.Optional[Any]:  # pragma: no cover
+        """Get the unsubscripted version of a type.
+        This supports generic types, Callable, Tuple, Union, Literal, Final and ClassVar.
+        Return None for unsupported types. Examples::
+            get_origin(Literal[42]) is Literal
+            get_origin(int) is None
+            get_origin(ClassVar[int]) is ClassVar
+            get_origin(Generic) is Generic
+            get_origin(Generic[T]) is Generic
+            get_origin(Union[T, int]) is Union
+            get_origin(List[Tuple[T, T]][int]) == list
+        """
+        if hasattr(tp, "__origin__"):
+            return tp.__origin__
+        if tp is typing.Generic:
+            return typing.Generic
+        return None
+
+
 try:
     BaseExceptionGroup = BaseExceptionGroup
     ExceptionGroup = ExceptionGroup  # pragma: no cover


Test patch:
diff --git a/hypothesis-python/tests/ghostwriter/recorded/addition_op_magic.txt b/hypothesis-python/tests/ghostwriter/recorded/addition_op_magic.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/addition_op_magic.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/addition_op_magic.txt
@@ -8,19 +8,19 @@ add_operands = st.floats()
 
 
 @given(a=add_operands, b=add_operands, c=add_operands)
-def test_associative_binary_operation_add(a, b, c):
+def test_associative_binary_operation_add(a: float, b: float, c) -> None:
     left = test_expected_output.add(a=a, b=test_expected_output.add(a=b, b=c))
     right = test_expected_output.add(a=test_expected_output.add(a=a, b=b), b=c)
     assert left == right, (left, right)
 
 
 @given(a=add_operands, b=add_operands)
-def test_commutative_binary_operation_add(a, b):
+def test_commutative_binary_operation_add(a: float, b: float) -> None:
     left = test_expected_output.add(a=a, b=b)
     right = test_expected_output.add(a=b, b=a)
     assert left == right, (left, right)
 
 
 @given(a=add_operands)
-def test_identity_binary_operation_add(a):
+def test_identity_binary_operation_add(a: float) -> None:
     assert a == test_expected_output.add(a=a, b=0.0)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/addition_op_multimagic.txt b/hypothesis-python/tests/ghostwriter/recorded/addition_op_multimagic.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/addition_op_multimagic.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/addition_op_multimagic.txt
@@ -8,7 +8,7 @@ from hypothesis import given, strategies as st
 
 
 @given(a=st.floats(), b=st.floats())
-def test_equivalent_add_add_add(a, b):
+def test_equivalent_add_add_add(a: float, b: float) -> None:
     result_0_add = _operator.add(a, b)
     result_1_add = numpy.add(a, b)
     result_2_add = test_expected_output.add(a=a, b=b)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/division_binop_error_handler.txt b/hypothesis-python/tests/ghostwriter/recorded/division_binop_error_handler.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/division_binop_error_handler.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/division_binop_error_handler.txt
@@ -8,19 +8,19 @@ divide_operands = st.integers()
 
 
 @given(a=divide_operands, b=divide_operands, c=divide_operands)
-def test_associative_binary_operation_divide(a, b, c):
+def test_associative_binary_operation_divide(a: int, b: int, c) -> None:
     left = test_expected_output.divide(a=a, b=test_expected_output.divide(a=b, b=c))
     right = test_expected_output.divide(a=test_expected_output.divide(a=a, b=b), b=c)
     assert left == right, (left, right)
 
 
 @given(a=divide_operands, b=divide_operands)
-def test_commutative_binary_operation_divide(a, b):
+def test_commutative_binary_operation_divide(a: int, b: int) -> None:
     left = test_expected_output.divide(a=a, b=b)
     right = test_expected_output.divide(a=b, b=a)
     assert left == right, (left, right)
 
 
 @given(a=divide_operands)
-def test_identity_binary_operation_divide(a):
+def test_identity_binary_operation_divide(a: int) -> None:
     assert a == test_expected_output.divide(a=a, b=1)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/division_fuzz_error_handler.txt b/hypothesis-python/tests/ghostwriter/recorded/division_fuzz_error_handler.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/division_fuzz_error_handler.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/division_fuzz_error_handler.txt
@@ -6,7 +6,7 @@ from hypothesis import given, reject, strategies as st
 
 
 @given(a=st.integers(), b=st.integers())
-def test_fuzz_divide(a, b):
+def test_fuzz_divide(a: int, b: int) -> None:
     try:
         test_expected_output.divide(a=a, b=b)
     except ZeroDivisionError:
diff --git a/hypothesis-python/tests/ghostwriter/recorded/division_operator_with_annotations.txt b/hypothesis-python/tests/ghostwriter/recorded/division_operator_with_annotations.txt
new file mode 100644
--- /dev/null
+++ b/hypothesis-python/tests/ghostwriter/recorded/division_operator_with_annotations.txt
@@ -0,0 +1,14 @@
+# This test code was written by the `hypothesis.extra.ghostwriter` module
+# and is provided under the Creative Commons Zero public domain dedication.
+
+import _operator
+from hypothesis import given, strategies as st
+
+# TODO: replace st.nothing() with an appropriate strategy
+
+truediv_operands = st.nothing()
+
+
+@given(a=truediv_operands)
+def test_identity_binary_operation_truediv(a) -> None:
+    assert a == _operator.truediv(a, "identity element here")
diff --git a/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_arithmeticerror_handler.txt b/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_arithmeticerror_handler.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_arithmeticerror_handler.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_arithmeticerror_handler.txt
@@ -7,7 +7,7 @@ from hypothesis import given, reject, strategies as st
 
 
 @given(a=st.integers(), b=st.integers())
-def test_roundtrip_divide_mul(a, b):
+def test_roundtrip_divide_mul(a: int, b: int) -> None:
     try:
         value0 = test_expected_output.divide(a=a, b=b)
         value1 = _operator.mul(value0, b)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_error_handler.txt b/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_error_handler.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_error_handler.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_error_handler.txt
@@ -7,7 +7,7 @@ from hypothesis import given, reject, strategies as st
 
 
 @given(a=st.integers(), b=st.integers())
-def test_roundtrip_divide_mul(a, b):
+def test_roundtrip_divide_mul(a: int, b: int) -> None:
     try:
         value0 = test_expected_output.divide(a=a, b=b)
     except ZeroDivisionError:
diff --git a/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_error_handler_without_annotations.txt b/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_error_handler_without_annotations.txt
new file mode 100644
--- /dev/null
+++ b/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_error_handler_without_annotations.txt
@@ -0,0 +1,16 @@
+# This test code was written by the `hypothesis.extra.ghostwriter` module
+# and is provided under the Creative Commons Zero public domain dedication.
+
+import _operator
+import test_expected_output
+from hypothesis import given, reject, strategies as st
+
+
+@given(a=st.integers(), b=st.integers())
+def test_roundtrip_divide_mul(a, b):
+    try:
+        value0 = test_expected_output.divide(a=a, b=b)
+    except ZeroDivisionError:
+        reject()
+    value1 = _operator.mul(value0, b)
+    assert a == value1, (a, value1)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_typeerror_handler.txt b/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_typeerror_handler.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_typeerror_handler.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/division_roundtrip_typeerror_handler.txt
@@ -7,7 +7,7 @@ from hypothesis import given, reject, strategies as st
 
 
 @given(a=st.integers(), b=st.integers())
-def test_roundtrip_divide_mul(a, b):
+def test_roundtrip_divide_mul(a: int, b: int) -> None:
     try:
         try:
             value0 = test_expected_output.divide(a=a, b=b)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/fuzz_classmethod.txt b/hypothesis-python/tests/ghostwriter/recorded/fuzz_classmethod.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/fuzz_classmethod.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/fuzz_classmethod.txt
@@ -6,5 +6,5 @@ from hypothesis import given, strategies as st
 
 
 @given(arg=st.integers())
-def test_fuzz_A_Class_a_classmethod(arg):
+def test_fuzz_A_Class_a_classmethod(arg: int) -> None:
     test_expected_output.A_Class.a_classmethod(arg=arg)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/fuzz_sorted_with_annotations.txt b/hypothesis-python/tests/ghostwriter/recorded/fuzz_sorted_with_annotations.txt
new file mode 100644
--- /dev/null
+++ b/hypothesis-python/tests/ghostwriter/recorded/fuzz_sorted_with_annotations.txt
@@ -0,0 +1,13 @@
+# This test code was written by the `hypothesis.extra.ghostwriter` module
+# and is provided under the Creative Commons Zero public domain dedication.
+
+from hypothesis import given, strategies as st
+
+
+@given(
+    iterable=st.one_of(st.iterables(st.integers()), st.iterables(st.text())),
+    key=st.none(),
+    reverse=st.booleans(),
+)
+def test_fuzz_sorted(iterable, key, reverse) -> None:
+    sorted(iterable, key=key, reverse=reverse)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/fuzz_staticmethod.txt b/hypothesis-python/tests/ghostwriter/recorded/fuzz_staticmethod.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/fuzz_staticmethod.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/fuzz_staticmethod.txt
@@ -6,5 +6,5 @@ from hypothesis import given, strategies as st
 
 
 @given(arg=st.integers())
-def test_fuzz_A_Class_a_staticmethod(arg):
+def test_fuzz_A_Class_a_staticmethod(arg: int) -> None:
     test_expected_output.A_Class.a_staticmethod(arg=arg)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/hypothesis_module_magic.txt b/hypothesis-python/tests/ghostwriter/recorded/hypothesis_module_magic.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/hypothesis_module_magic.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/hypothesis_module_magic.txt
@@ -1,7 +1,10 @@
 # This test code was written by the `hypothesis.extra.ghostwriter` module
 # and is provided under the Creative Commons Zero public domain dedication.
 
+import datetime
 import hypothesis
+import hypothesis.strategies
+import random
 import typing
 from hypothesis import given, settings, strategies as st
 from hypothesis.strategies._internal.strategies import Ex
@@ -9,12 +12,12 @@ from random import Random
 
 
 @given(condition=st.from_type(object))
-def test_fuzz_assume(condition):
+def test_fuzz_assume(condition: object) -> None:
     hypothesis.assume(condition=condition)
 
 
 @given(value=st.text())
-def test_fuzz_event(value):
+def test_fuzz_event(value: str) -> None:
     hypothesis.event(value=value)
 
 
@@ -25,7 +28,13 @@ def test_fuzz_event(value):
     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):
+def test_fuzz_find(
+    specifier: hypothesis.strategies.SearchStrategy,
+    condition: typing.Callable[[typing.Any], bool],
+    settings: typing.Union[hypothesis.settings, None],
+    random: typing.Union[random.Random, None],
+    database_key: typing.Union[bytes, None],
+) -> None:
     hypothesis.find(
         specifier=specifier,
         condition=condition,
@@ -36,22 +45,22 @@ def test_fuzz_find(specifier, condition, settings, random, database_key):
 
 
 @given(value=st.text())
-def test_fuzz_note(value):
+def test_fuzz_note(value: str) -> None:
     hypothesis.note(value=value)
 
 
 @given(r=st.builds(Random))
-def test_fuzz_register_random(r):
+def test_fuzz_register_random(r: random.Random) -> None:
     hypothesis.register_random(r=r)
 
 
 @given(version=st.text(), blob=st.binary())
-def test_fuzz_reproduce_failure(version, blob):
+def test_fuzz_reproduce_failure(version: str, blob: bytes) -> None:
     hypothesis.reproduce_failure(version=version, blob=blob)
 
 
 @given(seed=st.from_type(typing.Hashable))
-def test_fuzz_seed(seed):
+def test_fuzz_seed(seed: typing.Hashable) -> None:
     hypothesis.seed(seed=seed)
 
 
@@ -69,18 +78,18 @@ def test_fuzz_seed(seed):
     print_blob=st.just(not_set),
 )
 def test_fuzz_settings(
-    parent,
-    max_examples,
-    derandomize,
+    parent: typing.Union[hypothesis.settings, None],
+    max_examples: int,
+    derandomize: bool,
     database,
     verbosity,
     phases,
-    stateful_step_count,
-    report_multiple_bugs,
+    stateful_step_count: int,
+    report_multiple_bugs: bool,
     suppress_health_check,
-    deadline,
-    print_blob,
-):
+    deadline: typing.Union[int, float, datetime.timedelta, None],
+    print_blob: bool,
+) -> None:
     hypothesis.settings(
         parent=parent,
         max_examples=max_examples,
@@ -97,20 +106,22 @@ def test_fuzz_settings(
 
 
 @given(name=st.text())
-def test_fuzz_settings_get_profile(name):
+def test_fuzz_settings_get_profile(name: str) -> None:
     hypothesis.settings.get_profile(name=name)
 
 
 @given(name=st.text())
-def test_fuzz_settings_load_profile(name):
+def test_fuzz_settings_load_profile(name: str) -> None:
     hypothesis.settings.load_profile(name=name)
 
 
 @given(name=st.text(), parent=st.one_of(st.none(), st.builds(settings)))
-def test_fuzz_settings_register_profile(name, parent):
+def test_fuzz_settings_register_profile(
+    name: str, parent: typing.Union[hypothesis.settings, None]
+) -> None:
     hypothesis.settings.register_profile(name=name, parent=parent)
 
 
 @given(observation=st.one_of(st.floats(), st.integers()), label=st.text())
-def test_fuzz_target(observation, label):
+def test_fuzz_target(observation: typing.Union[int, float], label: str) -> None:
     hypothesis.target(observation=observation, label=label)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/magic_base64_roundtrip_with_annotations.txt b/hypothesis-python/tests/ghostwriter/recorded/magic_base64_roundtrip_with_annotations.txt
new file mode 100644
--- /dev/null
+++ b/hypothesis-python/tests/ghostwriter/recorded/magic_base64_roundtrip_with_annotations.txt
@@ -0,0 +1,14 @@
+# This test code was written by the `hypothesis.extra.ghostwriter` module
+# and is provided under the Creative Commons Zero public domain dedication.
+
+import base64
+from hypothesis import given, strategies as st
+
+# TODO: replace st.nothing() with an appropriate strategy
+
+
+@given(altchars=st.none(), s=st.nothing(), validate=st.booleans())
+def test_roundtrip_b64encode_b64decode(altchars, s, validate) -> None:
+    value0 = base64.b64encode(s=s, altchars=altchars)
+    value1 = base64.b64decode(s=value0, altchars=altchars, validate=validate)
+    assert s == value1, (s, value1)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/magic_class.txt b/hypothesis-python/tests/ghostwriter/recorded/magic_class.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/magic_class.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/magic_class.txt
@@ -6,15 +6,15 @@ from hypothesis import given, strategies as st
 
 
 @given()
-def test_fuzz_A_Class():
+def test_fuzz_A_Class() -> None:
     test_expected_output.A_Class()
 
 
 @given(arg=st.integers())
-def test_fuzz_A_Class_a_classmethod(arg):
+def test_fuzz_A_Class_a_classmethod(arg: int) -> None:
     test_expected_output.A_Class.a_classmethod(arg=arg)
 
 
 @given(arg=st.integers())
-def test_fuzz_A_Class_a_staticmethod(arg):
+def test_fuzz_A_Class_a_staticmethod(arg: int) -> None:
     test_expected_output.A_Class.a_staticmethod(arg=arg)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/sorted_self_equivalent_with_annotations.txt b/hypothesis-python/tests/ghostwriter/recorded/sorted_self_equivalent_with_annotations.txt
new file mode 100644
--- /dev/null
+++ b/hypothesis-python/tests/ghostwriter/recorded/sorted_self_equivalent_with_annotations.txt
@@ -0,0 +1,17 @@
+# This test code was written by the `hypothesis.extra.ghostwriter` module
+# and is provided under the Creative Commons Zero public domain dedication.
+
+from hypothesis import given, strategies as st
+
+
+@given(
+    iterable=st.one_of(st.iterables(st.integers()), st.iterables(st.text())),
+    key=st.none(),
+    reverse=st.booleans(),
+)
+def test_equivalent_sorted_sorted_sorted(iterable, key, reverse) -> None:
+    result_0_sorted = sorted(iterable, key=key, reverse=reverse)
+    result_1_sorted = sorted(iterable, key=key, reverse=reverse)
+    result_2_sorted = sorted(iterable, key=key, reverse=reverse)
+    assert result_0_sorted == result_1_sorted, (result_0_sorted, result_1_sorted)
+    assert result_0_sorted == result_2_sorted, (result_0_sorted, result_2_sorted)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/timsort_idempotent.txt b/hypothesis-python/tests/ghostwriter/recorded/timsort_idempotent.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/timsort_idempotent.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/timsort_idempotent.txt
@@ -2,11 +2,12 @@
 # and is provided under the Creative Commons Zero public domain dedication.
 
 import test_expected_output
+import typing
 from hypothesis import given, strategies as st
 
 
 @given(seq=st.one_of(st.binary(), st.lists(st.integers())))
-def test_idempotent_timsort(seq):
+def test_idempotent_timsort(seq: typing.Sequence[int]) -> None:
     result = test_expected_output.timsort(seq=seq)
     repeat = test_expected_output.timsort(seq=result)
     assert result == repeat, (result, repeat)
diff --git a/hypothesis-python/tests/ghostwriter/recorded/timsort_idempotent_asserts.txt b/hypothesis-python/tests/ghostwriter/recorded/timsort_idempotent_asserts.txt
--- a/hypothesis-python/tests/ghostwriter/recorded/timsort_idempotent_asserts.txt
+++ b/hypothesis-python/tests/ghostwriter/recorded/timsort_idempotent_asserts.txt
@@ -2,11 +2,12 @@
 # and is provided under the Creative Commons Zero public domain dedication.
 
 import test_expected_output
+import typing
 from hypothesis import given, reject, strategies as st
 
 
 @given(seq=st.one_of(st.binary(), st.lists(st.integers())))
-def test_idempotent_timsort(seq):
+def test_idempotent_timsort(seq: typing.Sequence[int]) -> None:
     try:
         result = test_expected_output.timsort(seq=seq)
         repeat = test_expected_output.timsort(seq=result)
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
@@ -88,12 +88,17 @@ def divide(a: int, b: int) -> float:
     "data",
     [
         ("fuzz_sorted", ghostwriter.fuzz(sorted)),
+        ("fuzz_sorted_with_annotations", ghostwriter.fuzz(sorted, annotate=True)),
         ("fuzz_with_docstring", ghostwriter.fuzz(with_docstring)),
         ("fuzz_classmethod", ghostwriter.fuzz(A_Class.a_classmethod)),
         ("fuzz_staticmethod", ghostwriter.fuzz(A_Class.a_staticmethod)),
         ("fuzz_ufunc", ghostwriter.fuzz(numpy.add)),
         ("magic_gufunc", ghostwriter.magic(numpy.matmul)),
         ("magic_base64_roundtrip", ghostwriter.magic(base64.b64encode)),
+        (
+            "magic_base64_roundtrip_with_annotations",
+            ghostwriter.magic(base64.b64encode, annotate=True),
+        ),
         ("re_compile", ghostwriter.fuzz(re.compile)),
         (
             "re_compile_except",
@@ -114,6 +119,10 @@ def divide(a: int, b: int) -> float:
         ),
         ("eval_equivalent", ghostwriter.equivalent(eval, ast.literal_eval)),
         ("sorted_self_equivalent", ghostwriter.equivalent(sorted, sorted, sorted)),
+        (
+            "sorted_self_equivalent_with_annotations",
+            ghostwriter.equivalent(sorted, sorted, sorted, annotate=True),
+        ),
         ("addition_op_magic", ghostwriter.magic(add)),
         ("addition_op_multimagic", ghostwriter.magic(add, operator.add, numpy.add)),
         ("division_fuzz_error_handler", ghostwriter.fuzz(divide)),
@@ -125,6 +134,10 @@ def divide(a: int, b: int) -> float:
             "division_roundtrip_error_handler",
             ghostwriter.roundtrip(divide, operator.mul),
         ),
+        (
+            "division_roundtrip_error_handler_without_annotations",
+            ghostwriter.roundtrip(divide, operator.mul, annotate=False),
+        ),
         (
             "division_roundtrip_arithmeticerror_handler",
             ghostwriter.roundtrip(divide, operator.mul, except_=ArithmeticError),
@@ -139,6 +152,12 @@ def divide(a: int, b: int) -> float:
                 operator.truediv, associative=False, commutative=False
             ),
         ),
+        (
+            "division_operator_with_annotations",
+            ghostwriter.binary_operation(
+                operator.truediv, associative=False, commutative=False, annotate=True
+            ),
+        ),
         (
             "multiplication_operator",
             ghostwriter.binary_operation(
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
@@ -315,7 +315,12 @@ class MyError(UnicodeDecodeError):
 )
 def test_exception_deduplication(exceptions, output):
     _, body = ghostwriter._make_test_body(
-        lambda: None, ghost="", test_body="pass", except_=exceptions, style="pytest"
+        lambda: None,
+        ghost="",
+        test_body="pass",
+        except_=exceptions,
+        style="pytest",
+        annotate=False,
     )
     assert f"except {output}:" in body
 
diff --git a/hypothesis-python/tests/ghostwriter/test_ghostwriter_cli.py b/hypothesis-python/tests/ghostwriter/test_ghostwriter_cli.py
--- a/hypothesis-python/tests/ghostwriter/test_ghostwriter_cli.py
+++ b/hypothesis-python/tests/ghostwriter/test_ghostwriter_cli.py
@@ -48,6 +48,9 @@
         ("hypothesis.errors.StopTest", lambda: fuzz(StopTest)),
         # Search for identity element does not print e.g. "You can use @seed ..."
         ("--binary-op operator.add", lambda: binary_operation(operator.add)),
+        # Annotations are passed correctly
+        ("sorted --annotate", lambda: fuzz(sorted, annotate=True)),
+        ("sorted --no-annotate", lambda: fuzz(sorted, annotate=False)),
     ],
 )
 def test_cli_python_equivalence(cli, code):