[WIP] Addition of UnitOperationWarning and strict kwarg to UnitRegistry

docs/usage.rst

@@ -1024,6 +1024,20 @@ When writing tests, it is convenient to use :mod:`unyt.testing`. In particular,
   >>> desired = actual.to("cm")
   >>> assert_allclose_units(actual, desired)
 
+Integrating :mod:`unyt` Into a Legacy Code with Non-Strict Mode
+---------------------------------------------------------------
+
+If using a custom :class:`UnitRegistry <unyt.unit_registry.UnitRegistry>`, it
+is possible to supply ``strict=False`` when initializing. This will change the behavior
+when an invalid operation is attempted. Instead of raising a :class:`UnitOperationError <unyt.exceptions.UnitOperationError>`,
+a :class:`UnitOperationWarning <unyt.exceptions.UnitOperationError>` will be
+provided instead, units will be stripped, and the operation will return a value
+without units.
+
+This behavior is not recommended in general since many of the important benefits
+of :mod:`unyt` are lost, but it can be useful for integrating :mod:`unyt` into
+legacy code which was not previously unit-aware so that unit support can be
+added gradually.
 
 Custom Unit Systems
 -------------------

unyt/array.py

@@ -125,6 +125,7 @@
     UnitConversionError,
     UnitsNotReducible,
     SymbolNotFoundError,
+    UnitOperationWarning,
 )
 from unyt.equivalencies import equivalence_registry
 from unyt._on_demand_imports import _astropy, _pint
@@ -140,6 +141,7 @@
 
 NULL_UNIT = Unit()
 POWER_SIGN_MAPPING = {multiply: 1, divide: -1}
+STRICT = True
 
 __doctest_requires__ = {
     ("unyt_array.from_pint", "unyt_array.to_pint"): ["pint"],
@@ -166,6 +168,15 @@ def _iterable(obj):
     return True
 
 
+def _unit_operation_error_raise_or_warn(ufunc, u0, u1, func, *inputs):
+    if u0.registry.strict and u1.registry.strict:  # True by default
+        raise UnitOperationError(ufunc, u0, u1)
+    else:
+        warnings.warn(UnitOperationWarning(ufunc, u0, u1))
+        unwrapped_inputs = [i.value if isinstance(i, unyt_array) else i for i in inputs]
+        return func(*unwrapped_inputs)
+
+
 @lru_cache(maxsize=128, typed=False)
 def _sqrt_unit(unit):
     return 1, unit ** 0.5
@@ -1759,12 +1770,16 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
             elif ufunc is power:
                 u1 = inp1
                 if inp0.shape != () and inp1.shape != ():
-                    raise UnitOperationError(ufunc, u0, u1)
+                    return _unit_operation_error_raise_or_warn(
+                        ufunc, u0, u1, func, *inputs
+                    )
                 if isinstance(u1, unyt_array):
                     if u1.units.is_dimensionless:
                         pass
                     else:
-                        raise UnitOperationError(ufunc, u0, u1.units)
+                        return _unit_operation_error_raise_or_warn(
+                            ufunc, u0, u1.units, func, *inputs
+                        )
                 if u1.shape == ():
                     u1 = float(u1)
                 else:
@@ -1814,9 +1829,13 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
                                         ret = bool(ret)
                                     return ret
                                 else:
-                                    raise UnitOperationError(ufunc, u0, u1)
+                                    return _unit_operation_error_raise_or_warn(
+                                        ufunc, u0, u1, func, *inputs
+                                    )
                         else:
-                            raise UnitOperationError(ufunc, u0, u1)
+                            return _unit_operation_error_raise_or_warn(
+                                ufunc, u0, u1, func, *inputs
+                            )
                     conv, offset = u1.get_conversion_factor(u0, inp1.dtype)
                     new_dtype = np.dtype("f" + str(inp1.dtype.itemsize))
                     conv = new_dtype.type(conv)

unyt/exceptions.py

@@ -47,6 +47,29 @@ def __str__(self):
         return err
 
 
+class UnitOperationWarning(UserWarning):
+    """A warning that is raised when unit operations are not allowed but
+    when running with a UnitRegistry with strict=False. In this case,
+    the operation is allowed to continue but with unit information stripped.
+    """
+
+    def __init__(self, operation, unit1, unit2=None):
+        self.operation = operation
+        self.unit1 = unit1
+        self.unit2 = unit2
+        UserWarning.__init__(self)
+
+    def __str__(self):
+        err = (
+            'The %s operator for unyt_arrays with units "%s" '
+            '(dimensions "%s") ' % (self.operation, self.unit1, self.unit1.dimensions)
+        )
+        if self.unit2 is not None:
+            err += 'and "%s" (dimensions "%s") ' % (self.unit2, self.unit2.dimensions)
+        err += "is not well defined. Performing operation without units instead."
+        return err
+
+
 class UnitConversionError(Exception):
     """An error raised when converting to a unit with different dimensions.
 

unyt/tests/test_unyt_array.py

@@ -57,6 +57,7 @@
     UnitOperationError,
     UnitParseError,
     UnitsNotReducible,
+    UnitOperationWarning,
 )
 from unyt.testing import assert_allclose_units, _process_warning
 from unyt.unit_symbols import cm, m, g, degree
@@ -2539,3 +2540,15 @@ def test_invalid_unit_quantity_from_string(s):
         match="Could not find unit symbol '{}' in the provided symbols.".format(un_str),
     ):
         unyt_quantity.from_string(s)
+
+
+def test_non_strict_registry():
+
+    reg = UnitRegistry(strict=False)
+
+    a1 = unyt_array([1, 2, 3], "m", registry=reg)
+    a2 = unyt_array([4, 5, 6], "kg", registry=reg)
+
+    with pytest.warns(UnitOperationWarning):
+        answer = operator.add(a1, a2)
+        assert_array_equal(answer, [5, 7, 9])

unyt/unit_registry.py

@@ -55,7 +55,9 @@ class UnitRegistry:
 
     _unit_system_id = None
 
-    def __init__(self, add_default_symbols=True, lut=None, unit_system=None):
+    def __init__(
+        self, add_default_symbols=True, lut=None, unit_system=None, strict=True
+    ):
         self._unit_object_cache = {}
         if lut:
             self.lut = lut
@@ -67,6 +69,12 @@ def __init__(self, add_default_symbols=True, lut=None, unit_system=None):
         if add_default_symbols:
             self.lut.update(default_unit_symbol_lut)
 
+        # This boolean determines whether to raise a UnitOperationError or
+        # strip units and provide a UnitOperationWarning if an invalid
+        # operation is attempted. The default is strict=True and this is
+        # strongly recommended.
+        self.strict = strict
+
     def __getitem__(self, key):
         try:
             ret = self.lut[str(key)]