diff options
Diffstat (limited to 'pint')
| -rw-r--r-- | pint/facets/numpy/quantity.py | 7 | ||||
| -rw-r--r-- | pint/facets/plain/quantity.py | 10 | ||||
| -rw-r--r-- | pint/facets/plain/registry.py | 6 | ||||
| -rw-r--r-- | pint/testsuite/test_numpy.py | 18 |
4 files changed, 36 insertions, 5 deletions
diff --git a/pint/facets/numpy/quantity.py b/pint/facets/numpy/quantity.py index 0d335cd..9aa55ce 100644 --- a/pint/facets/numpy/quantity.py +++ b/pint/facets/numpy/quantity.py @@ -244,7 +244,12 @@ class NumpyQuantity: def __setitem__(self, key, value): try: - if np.ma.is_masked(value) or math.isnan(value): + # If we're dealing with a masked single value or a nan, set it + if ( + isinstance(self._magnitude, np.ma.MaskedArray) + and np.ma.is_masked(value) + and getattr(value, "size", 0) == 1 + ) or math.isnan(value): self._magnitude[key] = value return except TypeError: diff --git a/pint/facets/plain/quantity.py b/pint/facets/plain/quantity.py index f4608c7..359e613 100644 --- a/pint/facets/plain/quantity.py +++ b/pint/facets/plain/quantity.py @@ -1363,6 +1363,14 @@ class PlainQuantity(PrettyIPython, SharedRegistryObject, Generic[_MagnitudeType] __rmatmul__ = __matmul__ + def _truedivide_cast_int(self, a, b): + t = self._REGISTRY.non_int_type + if isinstance(a, int): + a = t(a) + if isinstance(b, int): + b = t(b) + return operator.truediv(a, b) + def __itruediv__(self, other): if is_duck_array_type(type(self._magnitude)): return self._imul_div(other, operator.itruediv) @@ -1370,6 +1378,8 @@ class PlainQuantity(PrettyIPython, SharedRegistryObject, Generic[_MagnitudeType] return self._mul_div(other, operator.truediv) def __truediv__(self, other): + if isinstance(self.m, int) or isinstance(getattr(other, "m", None), int): + return self._mul_div(other, self._truedivide_cast_int, operator.truediv) return self._mul_div(other, operator.truediv) def __rtruediv__(self, other): diff --git a/pint/facets/plain/registry.py b/pint/facets/plain/registry.py index 255c7a5..0bf1545 100644 --- a/pint/facets/plain/registry.py +++ b/pint/facets/plain/registry.py @@ -1137,7 +1137,7 @@ class PlainRegistry(metaclass=RegistryMeta): token_text = token[1] if token_type == NAME: if token_text == "dimensionless": - return self.non_int_type("1") * self.dimensionless + return self.Quantity(1, self.dimensionless) elif token_text.lower() in ("inf", "infinity"): return self.non_int_type("inf") elif token_text.lower() == "nan": @@ -1146,7 +1146,7 @@ class PlainRegistry(metaclass=RegistryMeta): return self.Quantity(values[token_text]) else: return self.Quantity( - self.non_int_type("1"), + 1, self.UnitsContainer( {self.get_name(token_text, case_sensitive=case_sensitive): 1} ), @@ -1254,7 +1254,7 @@ class PlainRegistry(metaclass=RegistryMeta): ) if not input_string: - return self.Quantity(self.non_int_type("1")) + return self.Quantity(1) for p in self.preprocessors: input_string = p(input_string) diff --git a/pint/testsuite/test_numpy.py b/pint/testsuite/test_numpy.py index f9f1095..1e0b928 100644 --- a/pint/testsuite/test_numpy.py +++ b/pint/testsuite/test_numpy.py @@ -927,7 +927,7 @@ class TestNumpyUnclassified(TestNumpyMethods): q[:] = 1 * self.ureg.m helpers.assert_quantity_equal(q, [[1, 1], [1, 1]] * self.ureg.m) - # check and see that dimensionless num bers work correctly + # check and see that dimensionless numbers work correctly q = [0, 1, 2, 3] * self.ureg.dimensionless q[0] = 1 helpers.assert_quantity_equal(q, np.asarray([1, 1, 2, 3])) @@ -948,6 +948,22 @@ class TestNumpyUnclassified(TestNumpyMethods): assert not w assert q.mask[0] + def test_setitem_mixed_masked(self): + masked = np.ma.array( + [ + 1, + 2, + ], + mask=[True, False], + ) + q = self.Q_(np.ones(shape=(2,)), "m") + with pytest.raises(DimensionalityError): + q[:] = masked + + masked_q = self.Q_(masked, "mm") + q[:] = masked_q + helpers.assert_quantity_equal(q, [1.0, 0.002] * self.ureg.m) + def test_iterator(self): for q, v in zip(self.q.flatten(), [1, 2, 3, 4]): assert q == v * self.ureg.m |