Skip to content

Commit 7991cec

Browse files
committed
fix: Quantity repr/format crashes and wrong format-spec semantics
Fixes nine confirmed bugs from the string-formatting audit: - _format_value: catch OverflowError from jnp/np.asarray so repr/str no longer crash on Python ints exceeding int32 (e.g. 2**40). - __format__ '%' guard: use should_display_unit instead of is_unitless so named dimensionless units (radian) are rejected too, and test only the spec's type character so '%' as a fill char (e.g. '%>10') is allowed. - __format__ arrays: apply the format spec per element via np.array2string so 'e'/'g'/'%' keep their Python semantics instead of being treated as decimal rounding. - __format__ backend safety: degrade to str(self) for JAX tracers (no more crashes under jit), dask (no silent graph materialization), ndonnx (no opaque ufunc TypeError), and torch tensors with requires_grad. Each fix is pinned by a regression test in saiunit/_base_quantity_format_test.py.
1 parent adfe21d commit 7991cec

2 files changed

Lines changed: 233 additions & 16 deletions

File tree

saiunit/_base_quantity.py

Lines changed: 36 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
import functools
1919
import numbers
2020
import operator
21-
import re
2221
from collections.abc import Callable, Sequence
2322
from copy import deepcopy
2423
from typing import Any
@@ -1159,9 +1158,12 @@ def _format_value(self, precision: int | None = None) -> str:
11591158
# Prefer ``jnp.asarray`` when JAX is installed so the precision
11601159
# matches JAX's x32 default (test fixtures key off this); fall
11611160
# back to NumPy when JAX is unavailable.
1161+
# OverflowError: Python ints that exceed the backend integer
1162+
# dtype (e.g. 2**40 under JAX's x32 default). Keep the raw
1163+
# value; the fallback below prints it via ``str``.
11621164
try:
11631165
value = jnp.asarray(m) if HAS_JAX else np.asarray(m)
1164-
except TypeError:
1166+
except (TypeError, OverflowError):
11651167
value = m
11661168
else:
11671169
# cupy / torch / dask / ndonnx arrays: don't lift to JAX (that
@@ -1560,31 +1562,49 @@ def __str__(self) -> str:
15601562
def __format__(self, format_spec) -> str:
15611563
if not format_spec:
15621564
return str(self)
1563-
# Block '%' format on quantities with units — "50% mV" is meaningless
1564-
if '%' in format_spec and not self.unit.is_unitless:
1565+
# Block '%'-type format on quantities whose unit is displayed —
1566+
# "50% mV" (or "50% rad") is meaningless. The format type is always
1567+
# the last character of the spec, so a '%' used as a fill character
1568+
# (e.g. '%>10') is not rejected.
1569+
if format_spec.endswith('%') and self.unit.should_display_unit:
15651570
raise ValueError(
15661571
f"'%' format is not supported for Quantity with unit {str(self.unit)!r}. "
15671572
f"Convert to a dimensionless value first."
15681573
)
15691574
unit_str = str(self.unit)
15701575
show_unit = self.unit.should_display_unit
15711576
if self.shape == ():
1572-
formatted_value = format(self.mantissa, format_spec)
1577+
try:
1578+
formatted_value = format(self.mantissa, format_spec)
1579+
except TypeError:
1580+
# Mantissas that don't implement format specs (JAX tracers,
1581+
# lazy/symbolic backends): degrade to the plain display.
1582+
return str(self)
15731583
if not show_unit:
15741584
return formatted_value
15751585
return f"{formatted_value} {unit_str}"
15761586
else:
1577-
# Parse precision from standard format specs like .2f, .3e, .4g,
1578-
# 10.2f, +.2f, etc. Use a regex to extract the precision field.
1579-
m = re.match(r'^[^.]*\.(\d+)[feEgGn%]?$', format_spec)
1580-
if m is not None:
1581-
precision = int(m.group(1))
1582-
value = np.asarray(self.mantissa)
1583-
s = np.array_str(np.round(value, precision), precision=precision)
1584-
if not show_unit:
1585-
return s
1586-
return f"{s} {unit_str}"
1587-
return str(self)
1587+
# Apply the spec per element so 'e'/'g'/'%' types keep their
1588+
# Python semantics. Traced / lazy / symbolic mantissas cannot
1589+
# be converted to numpy without crashing or materializing —
1590+
# degrade to the plain display, exactly like ``repr`` does.
1591+
from saiunit._backend import is_dask_array, is_ndonnx_array
1592+
mantissa = self.mantissa
1593+
if _is_tracer(mantissa) or is_dask_array(mantissa) or is_ndonnx_array(mantissa):
1594+
return str(self)
1595+
try:
1596+
value = np.asarray(mantissa)
1597+
except (TypeError, ValueError, RuntimeError):
1598+
# e.g. torch tensors with requires_grad refuse __array__
1599+
return str(self)
1600+
try:
1601+
s = np.array2string(value, formatter={'all': lambda x: format(x, format_spec)})
1602+
except (TypeError, ValueError):
1603+
# Spec not applicable element-wise — fall back to str()
1604+
return str(self)
1605+
if not show_unit:
1606+
return s
1607+
return f"{s} {unit_str}"
15881608

15891609
def __iter__(self):
15901610
"""Solve the issue of DeviceArray.__iter__.
Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ==============================================================================
15+
16+
"""Regression tests from the Quantity string-formatting audit.
17+
18+
Each test below pins one confirmed bug in ``__repr__`` / ``__str__`` /
19+
``__format__`` / ``_format_value`` (saiunit/_base_quantity.py). The docstrings
20+
record the wrong behavior observed before the fix.
21+
"""
22+
23+
import jax
24+
import jax.numpy as jnp
25+
import numpy as np
26+
import pytest
27+
28+
import saiunit as u
29+
from saiunit import Quantity
30+
31+
32+
class TestReprLargePythonInt:
33+
def test_repr_does_not_crash_on_large_python_int(self):
34+
"""``_format_value`` promotes Python scalars via ``jnp.asarray`` and
35+
caught only TypeError. A Python int that does not fit int32 raises
36+
OverflowError, so ``repr``/``str`` crashed for values as small as
37+
2**40.
38+
39+
Before the fix: OverflowError: Python int 1099511627776 too large to
40+
convert to int32.
41+
"""
42+
q = Quantity(2 ** 40, unit=u.mV)
43+
r = repr(q)
44+
s = str(q)
45+
assert "1099511627776" in r
46+
assert "1099511627776" in s
47+
48+
49+
class TestPercentFormatGuard:
50+
def test_percent_rejected_for_named_dimensionless_unit(self):
51+
"""The '%' guard in ``__format__`` checked ``unit.is_unitless`` while
52+
display uses ``unit.should_display_unit``. Radian is unitless by
53+
magnitude but displays 'rad', so the guard let through exactly the
54+
nonsense it was written to block.
55+
56+
Before the fix: format(Quantity(0.5, unit=u.radian), '.1%')
57+
== '50.0% rad'.
58+
"""
59+
q = Quantity(0.5, unit=u.radian)
60+
with pytest.raises(ValueError):
61+
format(q, ".1%")
62+
63+
def test_percent_as_fill_character_is_not_rejected(self):
64+
"""The guard tested ``'%' in format_spec``, so '%' used as a
65+
legitimate fill character (e.g. '%>10') was rejected even though the
66+
format type is not '%'.
67+
68+
Before the fix: ValueError instead of '%%%%%%%0.5 mV'
69+
(plain ``format(0.5, '%>10')`` gives '%%%%%%%0.5').
70+
"""
71+
q = Quantity(0.5, unit=u.mV)
72+
assert format(q, "%>10") == "%%%%%%%0.5 mV"
73+
74+
def test_percent_array_keeps_percent_semantics(self):
75+
"""For non-scalar unitless quantities, '.1%' fell into the
76+
precision-regex path which treated '%' as a plain rounding spec: no
77+
x100 scaling and no '%' sign — silently wrong values. Scalars gave
78+
'50.0%'; arrays gave '[0.5 0.2]'.
79+
80+
Either honoring percent semantics per element or rejecting the spec
81+
for arrays is acceptable; silently emitting wrong numbers is not.
82+
"""
83+
q = Quantity(np.array([0.5, 0.25]))
84+
try:
85+
out = format(q, ".1%")
86+
except ValueError:
87+
return # explicit rejection is an acceptable contract
88+
assert "%" in out, f"percent semantics dropped: {out!r}"
89+
90+
91+
class TestFormatSpecArraySemantics:
92+
def test_scientific_spec_array_uses_scientific_notation(self):
93+
"""The array path extracted only the precision digit from '.1e' and
94+
applied decimal rounding (np.round(value, 1)), ignoring the 'e' type
95+
entirely. Scalars honor it ('1.2e+03'); arrays did not.
96+
97+
Before the fix: format(Quantity([1234.5], mV), '.1e')
98+
== '[1234.5] mV'.
99+
"""
100+
q = Quantity(np.array([1234.5]), unit=u.mV)
101+
out = format(q, ".1e")
102+
value_part = out.replace("mV", "")
103+
assert "e" in value_part, f"scientific notation dropped: {out!r}"
104+
105+
106+
class TestFormatSpecUnderJit:
107+
def test_format_spec_scalar_under_jit_does_not_crash(self):
108+
"""``__format__`` had no tracer guard (unlike ``_format_value``).
109+
For 0-d quantities it called ``format(self.mantissa, format_spec)``
110+
which raises on tracers.
111+
112+
Before the fix: TypeError: unsupported format string passed to
113+
DynamicJaxprTracer.__format__.
114+
"""
115+
captured = []
116+
117+
@jax.jit
118+
def f(x):
119+
q = Quantity(x, unit=u.mV)
120+
captured.append(f"{q:.2f}")
121+
return x
122+
123+
f(jnp.asarray(1.0))
124+
assert captured and isinstance(captured[0], str)
125+
126+
def test_format_spec_array_under_jit_does_not_crash(self):
127+
"""For non-scalar quantities ``__format__`` called
128+
``np.asarray(self.mantissa)`` which raises
129+
TracerArrayConversionError on traced arrays.
130+
"""
131+
captured = []
132+
133+
@jax.jit
134+
def f(x):
135+
q = Quantity(x, unit=u.mV)
136+
captured.append(f"{q:.2f}")
137+
return x
138+
139+
f(jnp.arange(3.0))
140+
assert captured and isinstance(captured[0], str)
141+
142+
143+
class TestFormatSpecLazyBackends:
144+
def test_format_spec_dask_does_not_materialize(self):
145+
"""``_format_value`` never materializes dask graphs and
146+
``np.asarray(Quantity)`` raises BackendError, but
147+
``__format__('.2f')`` called ``np.asarray(self.mantissa)`` directly,
148+
silently computing the whole graph.
149+
150+
Either a lazy-safe fallback or BackendError is acceptable; silent
151+
materialization is not.
152+
"""
153+
da = pytest.importorskip("dask.array")
154+
reads = []
155+
156+
def spy(block):
157+
if block.size: # ignore dask's zero-size meta-inference calls
158+
reads.append(block.size)
159+
return block
160+
161+
arr = da.from_array(np.array([1.0, 2.0]), chunks=1).map_blocks(spy)
162+
q = Quantity(arr, unit=u.mV)
163+
reads.clear()
164+
try:
165+
format(q, ".2f")
166+
except u.BackendError:
167+
pass # explicit rejection is an acceptable contract
168+
assert not reads, "format(q, '.2f') materialized the dask graph"
169+
170+
def test_format_spec_ndonnx_no_opaque_crash(self):
171+
"""``__format__('.2f')`` applied ``np.round`` to the ndonnx symbolic
172+
array, raising an opaque ufunc TypeError. ``repr``/``str`` degrade
173+
gracefully for the same quantity.
174+
175+
Before the fix: TypeError: loop of ufunc does not support argument 0
176+
of type Array which has no callable rint method.
177+
"""
178+
ndx = pytest.importorskip("ndonnx")
179+
q = Quantity(ndx.asarray(np.array([1.0, 2.0])), unit=u.mV)
180+
try:
181+
out = format(q, ".2f")
182+
except u.BackendError:
183+
return # explicit rejection is an acceptable contract
184+
assert isinstance(out, str)
185+
186+
def test_format_spec_torch_requires_grad_does_not_crash(self):
187+
"""``__format__('.2f')`` called ``np.asarray`` on the torch mantissa;
188+
tensors with ``requires_grad=True`` refuse the conversion.
189+
``repr``/``str`` work for the same quantity.
190+
191+
Before the fix: RuntimeError: Can't call numpy() on Tensor that
192+
requires grad.
193+
"""
194+
torch = pytest.importorskip("torch")
195+
q = Quantity(torch.tensor([1.5], requires_grad=True), unit=u.mV)
196+
out = format(q, ".2f")
197+
assert isinstance(out, str)

0 commit comments

Comments
 (0)