full separate numpy > 2.0 array wrap

zm711 · zm711 · commit b6efa33bb86d · 2024-07-24T13:24:21.000-04:00
diff --git a/quantities/quantity.py b/quantities/quantity.py
@@ -292,13 +292,18 @@ def __array_prepare__(self, obj, context=None):
 
     def __array_wrap__(self, obj, context=None, return_scalar=False):
         _np_version = tuple(map(int, np.__version__.split('.')))
+        # For NumPy < 2.0 we do old behavior
         if _np_version < (2, 0, 0):
             if not isinstance(obj, Quantity):
                 return self.__array_prepare__(obj, context)
             else:
                 return obj
+        # For NumPy > 2.0 we either do the prepare or the wrap
         else:
-            return super().__array_wrap__(obj, context, return_scalar)
+            if not isinstance(obj, Quantity):
+                return self.__array_prepare__(obj, context)
+            else:
+                return super().__array_wrap__(obj, context, return_scalar)
 
 
     @with_doc(np.ndarray.__add__)
