Merge pull request #3548 from ThunderKey/feature/add-annotations-to-g…

…hostwriter

Closes #3546

AUTHORS.rst

@@ -121,6 +121,7 @@ their individual contributions.
 * `Munir Abdinur <https://www.github.com/mabdinur>`_
 * `Nicholas Chammas <https://www.github.com/nchammas>`_
 * `Nick Anyos <https://www.github.com/NickAnyos>`_
+* `Nicolas Ganz <https://www.github.com/ThunderKey>`_
 * `Nikita Sobolev <https://github.com/sobolevn>`_ (mail@sobolevn.me)
 * `Oleg Höfling <https://github.com/hoefling>`_ (oleg.hoefling@gmail.com)
 * `Paul Ganssle <https://ganssle.io>`_ (paul@ganssle.io)

hypothesis-python/RELEASE.rst

@@ -0,0 +1,8 @@
+RELEASE_TYPE: minor
+
+The :doc:`Ghostwritter <ghostwriter>` will now include type annotations on tests
+for type-annotated code.  If you want to force this to happen (or not happen),
+pass a boolean to the new ``annotate=`` argument to the Python functions, or
+the ``--[no-]annotate`` CLI flag.
+
+Thanks to Nicolas Ganz for this new feature!

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: