fix: check_units result-callable receives units, not dimensions (#110)

The @check_units(result=callable) form passed each argument's *dimension*
to the result callable, but the returned value is validated by unit via
_check_unit/has_same_unit. has_same_unit calls get_unit() on the callable's
output; a Dimension collapses to Unit("1") there, so the comparison always
failed — even the documented example raised a spurious UnitMismatchError
("expected unit m^2 but got unit m^2").

Feed get_unit(a) instead so the callable produces a Unit, matching the
unit-based validation (check_dims correctly uses dimensions; assign_units
already uses units). Adds regression tests for the pass and mismatch paths.

Co-authored-by: Chaoming Wang <wcm.brainsim@gmail.com>

Author: chaoming0625

saiunit/_base_decorators.py

@@ -488,7 +488,10 @@ def new_f(*args, **kwds):
             result = f(*args, **kwds)
             if "result" in au:
                 if isinstance(au["result"], Callable) and au["result"] != bool:
-                    expected_result = au["result"](*[get_dim(a) for a in args])
+                    # Pass the argument *units* (not dimensions): the result is
+                    # validated by unit with ``_check_unit`` below, so a Dimension
+                    # here would always mismatch the returned Quantity's unit.
+                    expected_result = au["result"](*[get_unit(a) for a in args])
                 else:
                     expected_result = au["result"]
 

saiunit/_base_decorators_test.py

@@ -266,6 +266,27 @@ def f():
         with pytest.raises(UnitMismatchError):
             f()
 
+    def test_result_callable(self):
+        # Documented usage: the result unit is derived from the input unit.
+        # The callable must receive the argument *unit* (not its dimension),
+        # otherwise the unit comparison spuriously fails.
+        @check_units(result=lambda un: un ** 2)
+        def square(x):
+            return x ** 2
+
+        q = Quantity(3.0, unit=_metre)
+        result = square(q)
+        assert result.unit == _metre ** 2
+        assert jnp.allclose(result.mantissa, 9.0)
+
+    def test_result_callable_mismatch_raises(self):
+        @check_units(result=lambda un: un ** 2)
+        def bad(x):
+            return x  # wrong: returns metre, not metre**2
+
+        with pytest.raises(UnitMismatchError):
+            bad(Quantity(3.0, unit=_metre))
+
     def test_stores_metadata(self):
         @check_units(x=_metre, result=_metre)
         def f(x):