ENH: ensure yt.utils.funcs.ensure_numpy_array consistently returns a copy

Author: neutrinoceros

yt/_maintenance/numpy2_compat.py

@@ -1,9 +1,40 @@
 # avoid deprecation warnings in numpy >= 2.0
+from dataclasses import dataclass
+from importlib.metadata import version
+from typing import Literal, TypeAlias
 
 import numpy as np
+from packaging.version import Version
 
 if hasattr(np, "trapezoid"):
     # np.trapz is deprecated in numpy 2.0 in favor of np.trapezoid
     trapezoid = np.trapezoid
 else:
     trapezoid = np.trapz  # type: ignore # noqa: NPY201
+
+
+NUMPY_VERSION = Version(version("numpy"))
+
+# list numpy functions that gained a `copy` keyword argument at
+# some point after our oldest supported numpy version.
+CopyFunc: TypeAlias = Literal["asarray"]
+
+MIN_NUMPY_VERSION: dict[CopyFunc, Version] = {
+    "asarray": Version("2.0.0"),
+}
+
+
+@dataclass(frozen=True, slots=True)
+class CopyKwarg:
+    value: bool | None
+
+    def get(self, f: CopyFunc, /) -> dict[Literal["copy"], bool | None]:
+        if NUMPY_VERSION >= MIN_NUMPY_VERSION[f]:
+            return {"copy": self.value}
+        else:
+            return {}
+
+
+COPY_NONE = CopyKwarg(None)
+COPY_TRUE = CopyKwarg(True)
+COPY_FALSE = CopyKwarg(False)

yt/funcs.py

@@ -24,6 +24,7 @@
 
 from yt._maintenance.deprecation import issue_deprecation_warning
 from yt._maintenance.ipython_compat import IS_IPYTHON
+from yt._maintenance.numpy2_compat import COPY_TRUE
 from yt.config import ytcfg
 from yt.units import YTArray, YTQuantity
 from yt.utilities.exceptions import YTFieldNotFound, YTInvalidWidthError
@@ -86,15 +87,7 @@ def ensure_numpy_array(obj):
     This function ensures that *obj* is a numpy array. Typically used to
     convert scalar, list or tuple argument passed to functions using Cython.
     """
-    if isinstance(obj, np.ndarray):
-        if obj.shape == ():
-            return np.array([obj])
-        # We cast to ndarray to catch ndarray subclasses
-        return np.array(obj)
-    elif isinstance(obj, (list, tuple)):
-        return np.asarray(obj)
-    else:
-        return np.asarray([obj])
+    return np.atleast_1d(np.asarray(obj, **COPY_TRUE.get("asarray")))
 
 
 def read_struct(f, fmt):