Repository: hgrecco/pint

Issue:
synchronize Quantity.prod with np.prod(q)
Continuing from #867: #1087 fixed the calculated units for `prod` when the `axis` or `where` (not both) parameters are passed, but only for the ufunc and not for the method (also, right now this will calculate the units using input and output sizes, but instead we could just use `q.shape[axis]`).

I looked into implementing this but it seems this is somewhat complicated: the way it works right now, `__getattr__` will get the method from the magnitude and wrap it using `_numpy_method_wrap` which in turn will call `get_op_output_unit`. To make this work I think we'd need to move the units calculation of `_prod` into a new function that accepts the shape and the args (`axis` and `where`) and try to modify `get_op_output_unit` to accept functions in addition to names. @jthielen, thoughts?

Discussion:

Patch:
diff --git a/pint/numpy_func.py b/pint/numpy_func.py
index 5994461..aa03e18 100644
--- a/pint/numpy_func.py
+++ b/pint/numpy_func.py
@@ -416,7 +416,6 @@ matching_input_copy_units_output_ufuncs = [
 copy_units_output_ufuncs = ["ldexp", "fmod", "mod", "remainder"]
 op_units_output_ufuncs = {
     "var": "square",
-    "prod": "size",
     "multiply": "mul",
     "true_divide": "div",
     "divide": "div",
diff --git a/pint/quantity.py b/pint/quantity.py
index 1f18166..f5453d2 100644
--- a/pint/quantity.py
+++ b/pint/quantity.py
@@ -22,6 +22,7 @@ from typing import List
 from packaging import version
 
 from .compat import (
+    HAS_NUMPY_ARRAY_FUNCTION,
     NUMPY_VER,
     _to_magnitude,
     babel_parse,
@@ -112,6 +113,18 @@ def check_implemented(f):
     return wrapped
 
 
+def method_wraps(numpy_func):
+    if isinstance(numpy_func, str):
+        numpy_func = getattr(np, numpy_func, None)
+
+    def wrapper(func):
+        func.__wrapped__ = numpy_func
+
+        return func
+
+    return wrapper
+
+
 @contextlib.contextmanager
 def printoptions(*args, **kwargs):
     """Numpy printoptions context manager released with version 1.15.0
@@ -1696,6 +1709,23 @@ class Quantity(PrettyIPython, SharedRegistryObject):
 
         return np.dot(self, b)
 
+    @method_wraps("prod")
+    def prod(self, *args, **kwargs):
+        """ Return the product of quantity elements over a given axis
+
+        Wraps np.prod().
+        """
+        # TODO: remove after support for 1.16 has been dropped
+        if not HAS_NUMPY_ARRAY_FUNCTION:
+            raise NotImplementedError(
+                "prod is only defined for"
+                " numpy == 1.16 with NUMPY_ARRAY_FUNCTION_PROTOCOL enabled"
+                f" or for numpy >= 1.17 ({np.__version__} is installed)."
+                " Please try setting the NUMPY_ARRAY_FUNCTION_PROTOCOL environment variable"
+                " or updating your numpy version."
+            )
+        return np.prod(self, *args, **kwargs)
+
     def __ito_if_needed(self, to_units):
         if self.unitless and to_units == "radian":
             return


Test patch:
diff --git a/pint/testsuite/test_numpy.py b/pint/testsuite/test_numpy.py
index 9ecdcb1..4a6b7e2 100644
--- a/pint/testsuite/test_numpy.py
+++ b/pint/testsuite/test_numpy.py
@@ -283,8 +283,14 @@ class TestNumpyMathematicalFunctions(TestNumpyMethods):
 
     # Sums, products, differences
 
+    @helpers.requires_array_function_protocol()
     def test_prod(self):
-        self.assertEqual(self.q.prod(), 24 * self.ureg.m ** 4)
+        axis = 0
+        where = [[True, False], [True, True]]
+
+        self.assertQuantityEqual(self.q.prod(), 24 * self.ureg.m ** 4)
+        self.assertQuantityEqual(self.q.prod(axis=axis), [3, 8] * self.ureg.m ** 2)
+        self.assertQuantityEqual(self.q.prod(where=where), 12 * self.ureg.m ** 3)
 
     @helpers.requires_array_function_protocol()
     def test_prod_numpy_func(self):
@@ -783,9 +789,6 @@ class TestNumpyUnclassified(TestNumpyMethods):
         self.assertQuantityAlmostEqual(np.std(self.q), 1.11803 * self.ureg.m, rtol=1e-5)
         self.assertRaises(OffsetUnitCalculusError, np.std, self.q_temperature)
 
-    def test_prod(self):
-        self.assertEqual(self.q.prod(), 24 * self.ureg.m ** 4)
-
     def test_cumprod(self):
         self.assertRaises(DimensionalityError, self.q.cumprod)
         self.assertQuantityEqual((self.q / self.ureg.m).cumprod(), [1, 2, 6, 24])