Merge pull request #3570 from felixdivo/from_dtype-complex-bounds

Closes #3468

hypothesis-python/RELEASE.rst

@@ -0,0 +1,7 @@
+RELEASE_TYPE: minor
+
+This release allows for more precise generation of complex numbers using
+:func:`~hypothesis.extra.numpy.from_dtype`, by supporting the ``width``,
+``min_magnitude``, and ``min_magnitude`` arguments (:issue:`3468`).
+
+Thanks to Felix Divo for this feature!

hypothesis-python/src/hypothesis/extra/numpy.py

@@ -34,6 +34,7 @@
 from hypothesis.internal.coverage import check_function
 from hypothesis.internal.reflection import proxies
 from hypothesis.internal.validation import check_type
+from hypothesis.strategies._internal.numbers import Real
 from hypothesis.strategies._internal.strategies import T, check_strategy
 from hypothesis.strategies._internal.utils import defines_strategy
 
@@ -81,16 +82,18 @@ def from_dtype(
     allow_subnormal: Optional[bool] = None,
     exclude_min: Optional[bool] = None,
     exclude_max: Optional[bool] = None,
+    min_magnitude: Real = 0,
+    max_magnitude: Optional[Real] = None,
 ) -> st.SearchStrategy[Any]:
     """Creates a strategy which can generate any value of the given dtype.
 
-    Compatible ``**kwargs`` are passed to the inferred strategy function for
-    integers, floats, and strings.  This allows you to customise the min and max
-    values, control the length or contents of strings, or exclude non-finite
-    numbers.  This is particularly useful when kwargs are passed through from
+    Compatible parameters are passed to the inferred strategy function while
+    inapplicable ones are ignored.
+    This allows you, for example, to customise the min and max values,
+    control the length or contents of strings, or exclude non-finite
+    numbers. This is particularly useful when kwargs are passed through from
     :func:`arrays` which allow a variety of numeric dtypes, as it seamlessly
-    handles the ``width`` or representable bounds for you.  See :issue:`2552`
-    for more detail.
+    handles the ``width`` or representable bounds for you.
     """
     check_type(np.dtype, dtype, "dtype")
     kwargs = {k: v for k, v in locals().items() if k != "dtype" and v is not None}
@@ -136,15 +139,16 @@ def compat_kw(*args, **kw):
             ),
         )
     elif dtype.kind == "c":
-        # If anyone wants to add a `width` argument to `complex_numbers()`, we would
-        # accept a pull request and add passthrough support for magnitude bounds,
-        # but it's a low priority otherwise.
-        kws = compat_kw("allow_nan", "allow_infinity", "allow_subnormal")
-        if dtype.itemsize == 8:
-            float32 = st.floats(width=32, **kws)
-            result = st.builds(complex, float32, float32)
-        else:
-            result = st.complex_numbers(**kws)
+        result = st.complex_numbers(
+            width=min(8 * dtype.itemsize, 128),  # convert from bytes to bits
+            **compat_kw(
+                "min_magnitude",
+                "max_magnitude",
+                "allow_nan",
+                "allow_infinity",
+                "allow_subnormal",
+            ),
+        )
     elif dtype.kind in ("S", "a"):
         # Numpy strings are null-terminated; only allow round-trippable values.
         # `itemsize == 0` means 'fixed length determined at array creation'

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

@@ -1652,7 +1652,7 @@ def complex_numbers(
     of (around) floating-point epsilon, due to implementation via
     the system ``sqrt`` function.
 
-    The width argument specifies the maximum number of bits of precision
+    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.

hypothesis-python/tests/numpy/test_from_dtype.py

@@ -8,6 +8,8 @@
 # v. 2.0. If a copy of the MPL was not distributed with this file, You can
 # obtain one at https://mozilla.org/MPL/2.0/.
 
+import sys
+
 import numpy as np
 import pytest
 
@@ -197,9 +199,29 @@ def test_arrays_gives_useful_error_on_inconsistent_time_unit():
         (float, {"allow_nan": False}, lambda x: not np.isnan(x)),
         (float, {"allow_infinity": False}, lambda x: not np.isinf(x)),
         (float, {"allow_nan": False, "allow_infinity": False}, np.isfinite),
+        # Complex numbers: bounds and excluding nonfinites
         (complex, {"allow_nan": False}, lambda x: not np.isnan(x)),
         (complex, {"allow_infinity": False}, lambda x: not np.isinf(x)),
         (complex, {"allow_nan": False, "allow_infinity": False}, np.isfinite),
+        (
+            complex,
+            {"min_magnitude": 1e3},
+            lambda x: abs(x) >= 1e3 * (1 - sys.float_info.epsilon),
+        ),
+        (
+            complex,
+            {"max_magnitude": 1e2},
+            lambda x: abs(x) <= 1e2 * (1 + sys.float_info.epsilon),
+        ),
+        (
+            complex,
+            {"min_magnitude": 1, "max_magnitude": 1e6},
+            lambda x: (
+                (1 - sys.float_info.epsilon)
+                <= abs(x)
+                <= 1e6 * (1 + sys.float_info.epsilon)
+            ),
+        ),
         # Integer bounds, limited to the representable range
         ("int8", {"min_value": -1, "max_value": 1}, lambda x: -1 <= x <= 1),
         ("uint8", {"min_value": 1, "max_value": 2}, lambda x: 1 <= x <= 2),