From 2c2348ad08a28b2e7d8b7ec91279e26b14b9bcd8 Mon Sep 17 00:00:00 2001
From: Zac Hatfield-Dodds <zac.hatfield.dodds@gmail.com>
Date: Sun, 25 Feb 2024 01:16:44 -0800
Subject: [PATCH] Handle .map(np.array) and friends

---
 .../strategies/_internal/strategies.py        | 32 ++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py b/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py
index 9286083cd5..46d4005cdb 100644
--- a/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py
+++ b/hypothesis-python/src/hypothesis/strategies/_internal/strategies.py
@@ -852,7 +852,7 @@ def filter(self, condition: Callable[[Ex], Any]) -> "SearchStrategy[Ex]":
         ListStrategy = _list_strategy_type()
         if not isinstance(self.mapped_strategy, ListStrategy) or not (
             (isinstance(self.pack, type) and issubclass(self.pack, abc.Collection))
-            or self.pack is sorted
+            or self.pack in _collection_ish_functions()
         ):
             return super().filter(condition)
 
@@ -874,6 +874,36 @@ def _list_strategy_type():
     return ListStrategy
 
 
+def _collection_ish_functions():
+    funcs = [sorted]
+    if np := sys.modules.get("numpy"):
+        # c.f. https://numpy.org/doc/stable/reference/routines.array-creation.html
+        # Probably only `np.array` and `np.asarray` will be used in practice,
+        # but why should that stop us when we've already gone this far?
+        funcs += [
+            np.empty_like,
+            np.eye,
+            np.identity,
+            np.ones_like,
+            np.zeros_like,
+            np.array,
+            np.asarray,
+            np.asanyarray,
+            np.ascontiguousarray,
+            np.asmatrix,
+            np.copy,
+            np.rec.array,
+            np.rec.fromarrays,
+            np.rec.fromrecords,
+            np.diag,
+            # bonus undocumented functions from tab-completion:
+            np.asarray_chkfinite,
+            np.asfarray,
+            np.asfortranarray,
+        ]
+    return funcs
+
+
 filter_not_satisfied = UniqueIdentifier("filter not satisfied")