Add width to hypothesis.strategies.complex_numbers()

hypothesis-python/RELEASE.rst

@@ -0,0 +1,6 @@
+RELEASE_TYPE: minor
+
+This release adds a ``width`` parameter to :func:`hypothesis.strategies.complex_numbers`,
+analogously to :func:`hypothesis.strategies.floats`.
+
+Thanks to Felix Divo for the new feature!

hypothesis-python/src/hypothesis/strategies/_internal/core.py

@@ -1630,8 +1630,10 @@ def complex_numbers(
     allow_infinity: Optional[bool] = None,
     allow_nan: Optional[bool] = None,
     allow_subnormal: bool = True,
+    width: int = 128,
 ) -> SearchStrategy[complex]:
-    """Returns a strategy that generates complex numbers.
+    """Returns a strategy that generates :class:`~python:complex`
+    numbers.
 
     This strategy draws complex numbers with constrained magnitudes.
     The ``min_magnitude`` and ``max_magnitude`` parameters should be
@@ -1642,13 +1644,22 @@ def complex_numbers(
     is an error to enable ``allow_nan``.  If ``max_magnitude`` is finite,
     it is an error to enable ``allow_infinity``.
 
-    ``allow_subnormal`` is applied to each part of the complex number
-    separately, as for :func:`~hypothesis.strategies.floats`.
+    ``allow_infinity``, ``allow_nan``, and ``allow_subnormal`` are
+    applied to each part of the complex number separately, as for
+    :func:`~hypothesis.strategies.floats`.
 
     The magnitude constraints are respected up to a relative error
     of (around) floating-point epsilon, due to implementation via
     the system ``sqrt`` function.
 
+    The width argument specifies the maximum number of bits of precision
+    required to represent the entire generated complex number.
+    Valid values are 32, 64 or 128, which correspond to the real and imaginary
+    components each having width 16, 32 or 64, respectively.
+    Passing ``width=64`` will still use the builtin 128-bit
+    :class:`~python:complex` class, but always for values which can be
+    exactly represented as two 32-bit floats.
+
     Examples from this strategy shrink by shrinking their real and
     imaginary parts, as :func:`~hypothesis.strategies.floats`.
 
@@ -1677,15 +1688,23 @@ def complex_numbers(
             f"Cannot have allow_nan={allow_nan!r}, min_magnitude={min_magnitude!r} "
             f"max_magnitude={max_magnitude!r}"
         )
-
     check_type(bool, allow_subnormal, "allow_subnormal")
+    if width not in (32, 64, 128):
+        raise InvalidArgument(
+            f"width={width!r}, but must be 32, 64 or 128 (other complex dtypes "
+            "such as complex192 or complex256 are not supported)"
+            # For numpy, these types would be supported (but not by CPython):
+            # https://numpy.org/doc/stable/reference/arrays.scalars.html#complex-floating-point-types
+        )
+    component_width = width // 2
     allow_kw = {
         "allow_nan": allow_nan,
         "allow_infinity": allow_infinity,
         # If we have a nonzero normal min_magnitude and draw a zero imaginary part,
         # then allow_subnormal=True would be an error with the min_value to the floats()
         # strategy for the real part.  We therefore replace True with None.
         "allow_subnormal": None if allow_subnormal else allow_subnormal,
+        "width": component_width,
     }
 
     if min_magnitude == 0 and max_magnitude is None:

hypothesis-python/src/hypothesis/strategies/_internal/numbers.py

@@ -401,7 +401,7 @@ def floats(
 
     The width argument specifies the maximum number of bits of precision
     required to represent the generated float. Valid values are 16, 32, or 64.
-    Passing ``width=32`` will still use the builtin 64-bit ``float`` class,
+    Passing ``width=32`` will still use the builtin 64-bit :class:`~python:float` class,
     but always for values which can be exactly represented as a 32-bit float.
 
     The exclude_min and exclude_max argument can be used to generate numbers

hypothesis-python/tests/cover/test_direct_strategies.py

@@ -99,6 +99,8 @@ def fn_ktest(*fnkwargs):
     (ds.floats, {"exclude_max": True}),  # because max_value=None
     (ds.floats, {"min_value": 1.8, "width": 32}),
     (ds.floats, {"max_value": 1.8, "width": 32}),
+    (ds.complex_numbers, {"min_magnitude": 1.8, "width": 64}),
+    (ds.complex_numbers, {"max_magnitude": 1.8, "width": 64}),
     (ds.fractions, {"min_value": 2, "max_value": 1}),
     (ds.fractions, {"min_value": math.nan}),
     (ds.fractions, {"max_value": math.nan}),
@@ -149,6 +151,12 @@ def fn_ktest(*fnkwargs):
     (ds.complex_numbers, {"min_magnitude": 3, "max_magnitude": 2}),
     (ds.complex_numbers, {"max_magnitude": 2, "allow_infinity": True}),
     (ds.complex_numbers, {"max_magnitude": 2, "allow_nan": True}),
+    (ds.complex_numbers, {"width": None}),
+    # Conceivable mistake when misunderstanding width for individual component widths:
+    (ds.complex_numbers, {"width": 16}),
+    # Unsupported as of now:
+    (ds.complex_numbers, {"width": 196}),
+    (ds.complex_numbers, {"width": 256}),
     (ds.fixed_dictionaries, {"mapping": "fish"}),
     (ds.fixed_dictionaries, {"mapping": {1: "fish"}}),
     (ds.fixed_dictionaries, {"mapping": {}, "optional": "fish"}),
@@ -246,6 +254,9 @@ def test_validates_keyword_arguments(fn, kwargs):
     (ds.complex_numbers, {"allow_nan": False, "allow_infinity": True}),
     (ds.complex_numbers, {"allow_nan": False, "allow_infinity": False}),
     (ds.complex_numbers, {"max_magnitude": math.inf, "allow_infinity": True}),
+    (ds.complex_numbers, {"width": 32}),
+    (ds.complex_numbers, {"width": 64}),
+    (ds.complex_numbers, {"width": 128}),
     (ds.sampled_from, {"elements": [1]}),
     (ds.sampled_from, {"elements": [1, 2, 3]}),
     (ds.fixed_dictionaries, {"mapping": {1: ds.integers()}}),