Repository: hgrecco/pint

Issue:
numpy.concatenate Quantities and non-pint zeroes
This correctly crashes:
```python
>>> pint.Quantity(1, "kg") + 1
pint.errors.DimensionalityError: Cannot convert from 'kilogram' to 'dimensionless'
```
While this, also correctly, graciously detects that RHS is full of zeroes:
```python
>>> pint.Quantity([1, 2], "kg") + numpy.array([0, 0])
array([1, 2]) <Unit('kilogram')>
```
I would expect ``numpy.concatenate`` to behave in the same way as ``__add__``, but it doesn't:
```python
>>> numpy.concatenate([pint.Quantity([1, 2], "kg"), numpy.array([0])])
pint.errors.DimensionalityError: Cannot convert from 'dimensionless' to 'kilogram'
```
Expected output:
```python
array([1, 2, 0]) <Unit('kilogram')>
```

Discussion:

Patch:
diff --git a/CHANGES b/CHANGES
index ced7663..a49059d 100644
--- a/CHANGES
+++ b/CHANGES
@@ -12,8 +12,12 @@ Pint Changelog
   end_values (Issue #1026)
 - Reinstated support for pickle protocol 0 and 1, which is required by pytables
   (Issue #1036, Thanks Guido Imperiale)
-- NaN is now treated the same as zero in addition, equality, and disequality
-  (Issue #1051, Thanks Guido Imperiale)
+- Bare zeros and NaNs (not wrapped by Quantity) are now gracefully accepted by all numpy
+  operations; e.g. np.stack([Quantity([1, 2], "m"), [0, np.nan]) is now valid, whereas
+  np.stack([Quantity([1, 2], "m"), [3, 4]) will continue raising DimensionalityError.
+  (Issue #1050, Thanks Guido Imperiale)
+- NaN is now treated the same as zero in addition, subtraction, equality, and
+  disequality (Issue #1051, Thanks Guido Imperiale)
 
 
 0.11 (2020-02-19)
diff --git a/pint/numpy_func.py b/pint/numpy_func.py
index 210f967..c2ab623 100644
--- a/pint/numpy_func.py
+++ b/pint/numpy_func.py
@@ -48,14 +48,13 @@ def _is_sequence_with_quantity_elements(obj):
 
     Returns
     -------
-    bool
+    True if obj is a sequence and at least one element is a Quantity; False otherwise
     """
     return (
         iterable(obj)
         and sized(obj)
         and not isinstance(obj, str)
-        and len(obj) > 0
-        and all(_is_quantity(item) for item in obj)
+        and any(_is_quantity(item) for item in obj)
     )
 
 
@@ -67,42 +66,43 @@ def _get_first_input_units(args, kwargs=None):
         if _is_quantity(arg):
             return arg.units
         elif _is_sequence_with_quantity_elements(arg):
-            return arg[0].units
+            return next(arg_i.units for arg_i in arg if _is_quantity(arg_i))
+    raise TypeError("Expected at least one Quantity; found none")
 
 
 def convert_arg(arg, pre_calc_units):
     """Convert quantities and sequences of quantities to pre_calc_units and strip units.
 
-    Helper function for convert_to_consistent_units. pre_calc_units must be given as a pint
-    Unit or None.
-
+    Helper function for convert_to_consistent_units. pre_calc_units must be given as a
+    pint Unit or None.
     """
     if pre_calc_units is not None:
         if _is_quantity(arg):
             return arg.m_as(pre_calc_units)
         elif _is_sequence_with_quantity_elements(arg):
-            return [item.m_as(pre_calc_units) for item in arg]
+            return [convert_arg(item, pre_calc_units) for item in arg]
         elif arg is not None:
             if pre_calc_units.dimensionless:
                 return pre_calc_units._REGISTRY.Quantity(arg).m_as(pre_calc_units)
+            elif not _is_quantity(arg) and zero_or_nan(arg, True):
+                return arg
             else:
                 raise DimensionalityError("dimensionless", pre_calc_units)
-    else:
-        if _is_quantity(arg):
-            return arg.m
-        elif _is_sequence_with_quantity_elements(arg):
-            return [item.m for item in arg]
+    elif _is_quantity(arg):
+        return arg.m
+    elif _is_sequence_with_quantity_elements(arg):
+        return [convert_arg(item, pre_calc_units) for item in arg]
     return arg
 
 
 def convert_to_consistent_units(*args, pre_calc_units=None, **kwargs):
     """Prepare args and kwargs for wrapping by unit conversion and stripping.
 
-    If pre_calc_units is not None, takes the args and kwargs for a NumPy function and converts
-    any Quantity or Sequence of Quantities into the units of the first Quantiy/Sequence of
-    Quantities and returns the magnitudes. Other args/kwargs are treated as dimensionless
-    Quantities. If pre_calc_units is None, units are simply stripped.
-
+    If pre_calc_units is not None, takes the args and kwargs for a NumPy function and
+    converts any Quantity or Sequence of Quantities into the units of the first
+    Quantity/Sequence of Quantities and returns the magnitudes. Other args/kwargs are
+    treated as dimensionless Quantities. If pre_calc_units is None, units are simply
+    stripped.
     """
     return (
         tuple(convert_arg(arg, pre_calc_units=pre_calc_units) for arg in args),
@@ -484,29 +484,15 @@ def _power(x1, x2):
         return x2.__rpow__(x1)
 
 
-def _add_subtract_handle_non_quantity_zero(x1, x2):
-    # As in #121, #122, and #1051, if a value is 0 or NaN (but not a Quantity) do the
-    # operation without checking units. We do the calculation instead of just returning
-    # the same value to enforce any shape checking and type casting due to the
-    # operation.
-    if zero_or_nan(x1, True):
-        (x2,), output_wrap = unwrap_and_wrap_consistent_units(x2)
-    elif zero_or_nan(x2, True):
-        (x1,), output_wrap = unwrap_and_wrap_consistent_units(x1)
-    else:
-        (x1, x2), output_wrap = unwrap_and_wrap_consistent_units(x1, x2)
-    return x1, x2, output_wrap
-
-
 @implements("add", "ufunc")
 def _add(x1, x2, *args, **kwargs):
-    x1, x2, output_wrap = _add_subtract_handle_non_quantity_zero(x1, x2)
+    (x1, x2), output_wrap = unwrap_and_wrap_consistent_units(x1, x2)
     return output_wrap(np.add(x1, x2, *args, **kwargs))
 
 
 @implements("subtract", "ufunc")
 def _subtract(x1, x2, *args, **kwargs):
-    x1, x2, output_wrap = _add_subtract_handle_non_quantity_zero(x1, x2)
+    (x1, x2), output_wrap = unwrap_and_wrap_consistent_units(x1, x2)
     return output_wrap(np.subtract(x1, x2, *args, **kwargs))
 
 
@@ -550,26 +536,7 @@ def _interp(x, xp, fp, left=None, right=None, period=None):
 
 @implements("where", "function")
 def _where(condition, *args):
-    if (
-        len(args) == 2
-        and not _is_quantity(args[1])
-        and not iterable(args[1])
-        and (args[1] == 0 or np.isnan(args[1]))
-    ):
-        # Special case for y being bare zero or nan
-        (x,), output_wrap = unwrap_and_wrap_consistent_units(args[0])
-        args = x, args[1]
-    elif (
-        len(args) == 2
-        and not _is_quantity(args[0])
-        and not iterable(args[0])
-        and (args[0] == 0 or np.isnan(args[0]))
-    ):
-        # Special case for x being bare zero or nan
-        (y,), output_wrap = unwrap_and_wrap_consistent_units(args[1])
-        args = args[0], y
-    else:
-        args, output_wrap = unwrap_and_wrap_consistent_units(*args)
+    args, output_wrap = unwrap_and_wrap_consistent_units(*args)
     return output_wrap(np.where(condition, *args))
 
 
@@ -631,6 +598,8 @@ def _isin(element, test_elements, assume_unique=False, invert=False):
     elif _is_sequence_with_quantity_elements(test_elements):
         compatible_test_elements = []
         for test_element in test_elements:
+            if not _is_quantity(test_element):
+                pass
             try:
                 compatible_test_elements.append(test_element.m_as(element.units))
             except DimensionalityError:


Test patch:
diff --git a/pint/testsuite/test_numpy.py b/pint/testsuite/test_numpy.py
index 244d028..4e89172 100644
--- a/pint/testsuite/test_numpy.py
+++ b/pint/testsuite/test_numpy.py
@@ -29,6 +29,10 @@ class TestNumpyMethods(QuantityTestCase):
     def q_nan(self):
         return [[1, 2], [3, np.nan]] * self.ureg.m
 
+    @property
+    def q_zero_or_nan(self):
+        return [[0, 0], [0, np.nan]] * self.ureg.m
+
     @property
     def q_temperature(self):
         return self.Q_([[1, 2], [3, 4]], self.ureg.degC)
@@ -182,45 +186,60 @@ class TestNumpyArrayManipulation(TestNumpyMethods):
     # Changing number of dimensions
     # Joining arrays
     @helpers.requires_array_function_protocol()
-    def test_concatentate(self):
-        self.assertQuantityEqual(
-            np.concatenate([self.q] * 2),
-            self.Q_(np.concatenate([self.q.m] * 2), self.ureg.m),
-        )
-
-    @helpers.requires_array_function_protocol()
-    def test_stack(self):
-        self.assertQuantityEqual(
-            np.stack([self.q] * 2), self.Q_(np.stack([self.q.m] * 2), self.ureg.m)
-        )
-
-    @helpers.requires_array_function_protocol()
-    def test_column_stack(self):
-        self.assertQuantityEqual(np.column_stack([self.q[:, 0], self.q[:, 1]]), self.q)
-
-    @helpers.requires_array_function_protocol()
-    def test_dstack(self):
-        self.assertQuantityEqual(
-            np.dstack([self.q] * 2), self.Q_(np.dstack([self.q.m] * 2), self.ureg.m)
-        )
-
-    @helpers.requires_array_function_protocol()
-    def test_hstack(self):
-        self.assertQuantityEqual(
-            np.hstack([self.q] * 2), self.Q_(np.hstack([self.q.m] * 2), self.ureg.m)
-        )
-
-    @helpers.requires_array_function_protocol()
-    def test_vstack(self):
-        self.assertQuantityEqual(
-            np.vstack([self.q] * 2), self.Q_(np.vstack([self.q.m] * 2), self.ureg.m)
-        )
-
-    @helpers.requires_array_function_protocol()
-    def test_block(self):
-        self.assertQuantityEqual(
-            np.block([self.q[0, :], self.q[1, :]]), self.Q_([1, 2, 3, 4], self.ureg.m)
-        )
+    def test_concat_stack(self):
+        for func in (np.concatenate, np.stack, np.hstack, np.vstack, np.dstack):
+            with self.subTest(func=func):
+                self.assertQuantityEqual(
+                    func([self.q] * 2), self.Q_(func([self.q.m] * 2), self.ureg.m),
+                )
+                # One or more of the args is a bare array full of zeros or NaNs
+                self.assertQuantityEqual(
+                    func([self.q_zero_or_nan.m, self.q]),
+                    self.Q_(func([self.q_zero_or_nan.m, self.q.m]), self.ureg.m),
+                )
+                # One or more of the args is a bare array with at least one non-zero,
+                # non-NaN element
+                nz = self.q_zero_or_nan
+                nz.m[0, 0] = 1
+                with self.assertRaises(DimensionalityError):
+                    func([nz.m, self.q])
+
+    @helpers.requires_array_function_protocol()
+    def test_block_column_stack(self):
+        for func in (np.block, np.column_stack):
+            with self.subTest(func=func):
+
+                self.assertQuantityEqual(
+                    func([self.q[:, 0], self.q[:, 1]]),
+                    self.Q_(func([self.q[:, 0].m, self.q[:, 1].m]), self.ureg.m),
+                )
+
+                # One or more of the args is a bare array full of zeros or NaNs
+                self.assertQuantityEqual(
+                    func(
+                        [
+                            self.q_zero_or_nan[:, 0].m,
+                            self.q[:, 0],
+                            self.q_zero_or_nan[:, 1].m,
+                        ]
+                    ),
+                    self.Q_(
+                        func(
+                            [
+                                self.q_zero_or_nan[:, 0].m,
+                                self.q[:, 0].m,
+                                self.q_zero_or_nan[:, 1].m,
+                            ]
+                        ),
+                        self.ureg.m,
+                    ),
+                )
+                # One or more of the args is a bare array with at least one non-zero,
+                # non-NaN element
+                nz = self.q_zero_or_nan
+                nz.m[0, 0] = 1
+                with self.assertRaises(DimensionalityError):
+                    func([nz[:, 0].m, self.q[:, 0]])
 
     @helpers.requires_array_function_protocol()
     def test_append(self):
diff --git a/pint/testsuite/test_numpy_func.py b/pint/testsuite/test_numpy_func.py
index fb36c74..3a9453d 100644
--- a/pint/testsuite/test_numpy_func.py
+++ b/pint/testsuite/test_numpy_func.py
@@ -55,7 +55,8 @@ class TestNumPyFuncUtils(TestNumpyMethods):
             )
         )
         self.assertTrue(_is_sequence_with_quantity_elements(np.arange(4) * self.ureg.m))
-        self.assertFalse(_is_sequence_with_quantity_elements((self.Q_(0), True)))
+        self.assertTrue(_is_sequence_with_quantity_elements((self.Q_(0), 0)))
+        self.assertTrue(_is_sequence_with_quantity_elements((0, self.Q_(0))))
         self.assertFalse(_is_sequence_with_quantity_elements([1, 3, 5]))
         self.assertFalse(_is_sequence_with_quantity_elements(9 * self.ureg.m))
         self.assertFalse(_is_sequence_with_quantity_elements(np.arange(4)))