Allow wrapping astropy.units.Quantity

pydata · Nov 3, 2024 · 37d3510 · 37d3510
1 parent 5fd50c5
commit 37d3510
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/xarray/core/variable.py b/xarray/core/variable.py
@@ -320,10 +320,17 @@ def convert_non_numpy_type(data):
         else:
             data = np.asarray(data)
 
+    _is_array_like = isinstance(data, np.ndarray | np.generic)
+    _is_nep18 = hasattr(data, "__array_function__")
+    _has_array_api = hasattr(data, "__array_namespace__")
+    _has_unit = hasattr(data, "_unit")
+
+    # Allow `astropy.units.Quantity`
+    if _is_array_like and (_is_nep18 or _has_array_api) and _has_unit:
+        return cast("T_DuckArray", data)
+
     # immediately return array-like types except `numpy.ndarray` subclasses and `numpy` scalars
-    if not isinstance(data, np.ndarray | np.generic) and (
-        hasattr(data, "__array_function__") or hasattr(data, "__array_namespace__")
-    ):
+    if not _is_array_like and (_is_nep18 or _has_array_api):
         return cast("T_DuckArray", data)
 
     # validate whether the data is valid data types. Also, explicitly cast `numpy`