quantity checking functions

pint/numpy_util.py

@@ -0,0 +1,97 @@
+from .util import iterable, sized
+from .compat import np
+
+
+def is_quantity_with_scalar_magnitude(obj):
+    """Test for Quantity with scalar magnitude.
+
+    Parameters
+    ----------
+    obj : object
+
+
+    Returns
+    -------
+    True if obj is a Quantity with a scalar magnitude; False otherwise
+    """
+    return is_quantity(obj) and not iterable(obj._magnitude)
+
+
+def is_quantity_with_sequence_magnitude(obj):
+    """Test for Quantity with sequence magnitude.
+
+    Parameters
+    ----------
+    obj : object
+
+
+    Returns
+    -------
+    True if obj is a Quantity with a sequence magnitude; False otherwise
+
+    Examples
+    --------
+
+    >>> is_quantity_with_sequence_magnitude([1, 2, 3])
+    False
+
+    >>> is_quantity_with_sequence_magnitude([1, Q_(2, 'm'), 3])
+    False
+
+    >>> is_quantity_with_sequence_magnitude(Q_([1, 2, 3], 'm'))
+    True
+    """
+    return is_quantity(obj) and iterable(obj._magnitude)
+
+
+def is_sequence_with_quantity_elements(obj):
+    """Test for sequences of quantities.
+
+    Parameters
+    ----------
+    obj : object
+
+
+    Returns
+    -------
+    True if obj is a sequence and at least one element is a Quantity; False otherwise
+
+    Examples
+    --------
+
+    >>> is_sequence_with_quantity_elements([1, 2, 3])
+    False
+
+    >>> is_sequence_with_quantity_elements([1, Q_(2, 'm'), 3])
+    True
+
+    >>> is_sequence_with_quantity_elements(Q_([1, 2, 3], 'm'))
+    True
+    """
+    if np is not None and isinstance(obj, np.ndarray) and not obj.dtype.hasobject:
+        # If obj is a numpy array, avoid looping on all elements
+        # if dtype does not have objects
+        return False
+    return (
+        iterable(obj)
+        and sized(obj)
+        and not isinstance(obj, str)
+        and any(is_quantity(item) for item in obj)
+    )
+
+
+def is_quantity(obj):
+    """Test for _units and _magnitude attrs.
+
+    This is done in place of isinstance(Quantity, arg), which would cause a circular import.
+
+    Parameters
+    ----------
+    obj : Object
+
+
+    Returns
+    -------
+    bool
+    """
+    return hasattr(obj, "_units") and hasattr(obj, "_magnitude")

pint/testsuite/test_numpy_util.py

@@ -0,0 +1,81 @@
+from pint.numpy_util import (
+    is_quantity,
+    is_quantity_with_scalar_magnitude,
+    is_quantity_with_sequence_magnitude,
+    is_sequence_with_quantity_elements,
+)
+from pint.compat import np
+import pytest
+from pint import Quantity as Q_
+
+
+@pytest.mark.parametrize(
+    "obj,result",
+    [
+        (Q_(1, "m"), True),
+        (Q_(np.nan, "m"), True),
+        (Q_([1, 2], "m"), True),
+        (Q_([1, np.nan], "m"), True),
+        (Q_(np.array([1, 2]), "m"), True),
+        (Q_(np.array([1, np.nan]), "m"), True),
+        (np.array([Q_(1, "m"), Q_(2, "m")], dtype="object"), False),
+        (np.array([Q_(1, "m"), Q_(np.nan, "m")], dtype="object"), False),
+        (np.array([Q_(1, "m"), np.nan], dtype="object"), False),
+    ],
+)
+def test_is_quantity(obj, result):
+    assert is_quantity(obj) == result
+
+
+@pytest.mark.parametrize(
+    "obj,result",
+    [
+        (Q_(1, "m"), True),
+        (Q_(np.nan, "m"), True),
+        (Q_([1, 2], "m"), False),
+        (Q_([1, np.nan], "m"), False),
+        (Q_(np.array([1, 2]), "m"), False),
+        (Q_(np.array([1, np.nan]), "m"), False),
+        (np.array([Q_(1, "m"), Q_(2, "m")], dtype="object"), False),
+        (np.array([Q_(1, "m"), Q_(np.nan, "m")], dtype="object"), False),
+        (np.array([Q_(1, "m"), np.nan], dtype="object"), False),
+    ],
+)
+def test_is_quantity_with_scalar_magnitude(obj, result):
+    assert is_quantity_with_scalar_magnitude(obj) == result
+
+
+@pytest.mark.parametrize(
+    "obj,result",
+    [
+        (Q_(1, "m"), False),
+        (Q_(np.nan, "m"), False),
+        (Q_([1, 2], "m"), True),
+        (Q_([1, np.nan], "m"), True),
+        (Q_(np.array([1, 2]), "m"), True),
+        (Q_(np.array([1, np.nan]), "m"), True),
+        (np.array([Q_(1, "m"), Q_(2, "m")], dtype="object"), False),
+        (np.array([Q_(1, "m"), Q_(np.nan, "m")], dtype="object"), False),
+        (np.array([Q_(1, "m"), np.nan], dtype="object"), False),
+    ],
+)
+def test_is_quantity_with_sequence_magnitude(obj, result):
+    assert is_quantity_with_sequence_magnitude(obj) == result
+
+
+@pytest.mark.parametrize(
+    "obj,result",
+    [
+        (Q_(1, "m"), False),
+        (Q_(np.nan, "m"), False),
+        (Q_([1, 2], "m"), True),
+        (Q_([1, np.nan], "m"), True),
+        (Q_(np.array([1, 2]), "m"), True),
+        (Q_(np.array([1, np.nan]), "m"), True),
+        (np.array([Q_(1, "m"), Q_(2, "m")], dtype="object"), True),
+        (np.array([Q_(1, "m"), Q_(np.nan, "m")], dtype="object"), True),
+        (np.array([Q_(1, "m"), np.nan], dtype="object"), True),
+    ],
+)
+def test_is_sequence_with_quantity_elements(obj, result):
+    assert is_sequence_with_quantity_elements(obj) == result