Skip to content

Commit 857a8e1

Browse files
authored
implement the inline repr (#22)
* implement the inline repr * use unit abbreviations * fall back to the repr for non-ndarray magnitudes * add tests for the inline repr * vendor format_array_flat * use maybe_truncate instead of slicing * add a test where the calculation of remaining chars becomes negative * configure pytest to use the xunit2 junit family * fix the pytest configuration
1 parent ca252c8 commit 857a8e1

File tree

3 files changed

+207
-2
lines changed

3 files changed

+207
-2
lines changed

pint_xarray/__init__.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,18 @@
33
except ImportError:
44
from importlib_metadata import version
55

6-
from . import testing # noqa
7-
from .accessors import PintDataArrayAccessor, PintDatasetAccessor # noqa
6+
import pint
7+
8+
from . import testing # noqa: F401
9+
from . import formatting
10+
from .accessors import PintDataArrayAccessor, PintDatasetAccessor # noqa: F401
811

912
try:
1013
__version__ = version("pint-xarray")
1114
except Exception:
1215
# Local copy or not installed with setuptools.
1316
# Disable minimum version checks on downstream libraries.
1417
__version__ = "999"
18+
19+
20+
pint.Quantity._repr_inline_ = formatting.inline_repr

pint_xarray/formatting.py

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
from itertools import zip_longest
2+
3+
import numpy as np
4+
5+
6+
# vendored from xarray.core.formatting
7+
# https://github.com/pydata/xarray/blob/v0.16.0/xarray/core/formatting.py#L18-216
8+
def pretty_print(x, numchars: int):
9+
"""Given an object `x`, call `str(x)` and format the returned string so
10+
that it is numchars long, padding with trailing spaces or truncating with
11+
ellipses as necessary
12+
"""
13+
s = maybe_truncate(x, numchars)
14+
return s + " " * max(numchars - len(s), 0)
15+
16+
17+
# vendored from xarray.core.formatting
18+
def maybe_truncate(obj, maxlen=500):
19+
s = str(obj)
20+
if len(s) > maxlen:
21+
s = s[: (maxlen - 3)] + "..."
22+
return s
23+
24+
25+
# vendored from xarray.core.formatting
26+
def wrap_indent(text, start="", length=None):
27+
if length is None:
28+
length = len(start)
29+
indent = "\n" + " " * length
30+
return start + indent.join(x for x in text.splitlines())
31+
32+
33+
# vendored from xarray.core.formatting
34+
def _get_indexer_at_least_n_items(shape, n_desired, from_end):
35+
assert 0 < n_desired <= np.prod(shape)
36+
cum_items = np.cumprod(shape[::-1])
37+
n_steps = np.argmax(cum_items >= n_desired)
38+
stop = int(np.ceil(float(n_desired) / np.r_[1, cum_items][n_steps]))
39+
indexer = (
40+
((-1 if from_end else 0),) * (len(shape) - 1 - n_steps)
41+
+ ((slice(-stop, None) if from_end else slice(stop)),)
42+
+ (slice(None),) * n_steps
43+
)
44+
return indexer
45+
46+
47+
# vendored from xarray.core.formatting
48+
def first_n_items(array, n_desired):
49+
"""Returns the first n_desired items of an array"""
50+
# Unfortunately, we can't just do array.flat[:n_desired] here because it
51+
# might not be a numpy.ndarray. Moreover, access to elements of the array
52+
# could be very expensive (e.g. if it's only available over DAP), so go out
53+
# of our way to get them in a single call to __getitem__ using only slices.
54+
if n_desired < 1:
55+
raise ValueError("must request at least one item")
56+
57+
if array.size == 0:
58+
# work around for https://github.com/numpy/numpy/issues/5195
59+
return []
60+
61+
if n_desired < array.size:
62+
indexer = _get_indexer_at_least_n_items(array.shape, n_desired, from_end=False)
63+
array = array[indexer]
64+
return np.asarray(array).flat[:n_desired]
65+
66+
67+
# vendored from xarray.core.formatting
68+
def last_n_items(array, n_desired):
69+
"""Returns the last n_desired items of an array"""
70+
# Unfortunately, we can't just do array.flat[-n_desired:] here because it
71+
# might not be a numpy.ndarray. Moreover, access to elements of the array
72+
# could be very expensive (e.g. if it's only available over DAP), so go out
73+
# of our way to get them in a single call to __getitem__ using only slices.
74+
if (n_desired == 0) or (array.size == 0):
75+
return []
76+
77+
if n_desired < array.size:
78+
indexer = _get_indexer_at_least_n_items(array.shape, n_desired, from_end=True)
79+
array = array[indexer]
80+
return np.asarray(array).flat[-n_desired:]
81+
82+
83+
# vendored from xarray.core.formatting
84+
def last_item(array):
85+
"""Returns the last item of an array in a list or an empty list."""
86+
if array.size == 0:
87+
# work around for https://github.com/numpy/numpy/issues/5195
88+
return []
89+
90+
indexer = (slice(-1, None),) * array.ndim
91+
return np.ravel(np.asarray(array[indexer])).tolist()
92+
93+
94+
# based on xarray.core.formatting.format_item
95+
def format_item(x, quote_strings=True):
96+
"""Returns a succinct summary of an object as a string"""
97+
if isinstance(x, (str, bytes)):
98+
return repr(x) if quote_strings else x
99+
elif isinstance(x, (float, np.float_)):
100+
return f"{x:.4}"
101+
else:
102+
return str(x)
103+
104+
105+
# based on xarray.core.formatting.format_item
106+
def format_items(x):
107+
"""Returns a succinct summaries of all items in a sequence as strings"""
108+
x = np.asarray(x)
109+
formatted = [format_item(xi) for xi in x]
110+
return formatted
111+
112+
113+
# vendored from xarray.core.formatting
114+
def format_array_flat(array, max_width: int):
115+
"""Return a formatted string for as many items in the flattened version of
116+
array that will fit within max_width characters.
117+
"""
118+
# every item will take up at least two characters, but we always want to
119+
# print at least first and last items
120+
max_possibly_relevant = min(
121+
max(array.size, 1), max(int(np.ceil(max_width / 2.0)), 2)
122+
)
123+
relevant_front_items = format_items(
124+
first_n_items(array, (max_possibly_relevant + 1) // 2)
125+
)
126+
relevant_back_items = format_items(last_n_items(array, max_possibly_relevant // 2))
127+
# interleave relevant front and back items:
128+
# [a, b, c] and [y, z] -> [a, z, b, y, c]
129+
relevant_items = sum(
130+
zip_longest(relevant_front_items, reversed(relevant_back_items)), ()
131+
)[:max_possibly_relevant]
132+
133+
cum_len = np.cumsum([len(s) + 1 for s in relevant_items]) - 1
134+
if (array.size > 2) and (
135+
(max_possibly_relevant < array.size) or (cum_len > max_width).any()
136+
):
137+
padding = " ... "
138+
count = min(
139+
array.size, max(np.argmax(cum_len + len(padding) - 1 > max_width), 2)
140+
)
141+
else:
142+
count = array.size
143+
padding = "" if (count <= 1) else " "
144+
145+
num_front = (count + 1) // 2
146+
num_back = count - num_front
147+
# note that num_back is 0 <--> array.size is 0 or 1
148+
# <--> relevant_back_items is []
149+
pprint_str = "".join(
150+
[
151+
" ".join(relevant_front_items[:num_front]),
152+
padding,
153+
" ".join(relevant_back_items[-num_back:]),
154+
]
155+
)
156+
157+
# As a final check, if it's still too long even with the limit in values,
158+
# replace the end with an ellipsis
159+
# NB: this will still returns a full 3-character ellipsis when max_width < 3
160+
if len(pprint_str) > max_width:
161+
pprint_str = pprint_str[: max(max_width - 3, 0)] + "..."
162+
163+
return pprint_str
164+
165+
166+
def inline_repr(quantity, max_width):
167+
magnitude = quantity.magnitude
168+
units = quantity.units
169+
170+
units_repr = f"{units:~P}"
171+
if isinstance(magnitude, np.ndarray):
172+
data_repr = format_array_flat(magnitude, max_width - len(units_repr) - 3)
173+
else:
174+
data_repr = maybe_truncate(repr(magnitude), max_width - len(units_repr) - 3)
175+
176+
return f"[{units_repr}] {data_repr}"

pint_xarray/tests/test_formatting.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
import pint
2+
import pytest
3+
4+
# only need to register _repr_inline_
5+
import pint_xarray # noqa: F401
6+
7+
unit_registry = pint.UnitRegistry(force_ndarray_like=True)
8+
9+
10+
@pytest.mark.parametrize(
11+
("length", "expected"),
12+
(
13+
(40, "[N] 7.1 5.4 9.8 21.4 15.3"),
14+
(20, "[N] 7.1 5.4 ... 15.3"),
15+
(10, "[N] 7.1..."),
16+
(7, "[N] ..."),
17+
(3, "[N] ..."),
18+
),
19+
)
20+
def test_inline_repr(length, expected):
21+
quantity = unit_registry.Quantity([7.1, 5.4, 9.8, 21.4, 15.3], "N")
22+
23+
assert quantity._repr_inline_(length) == expected

0 commit comments

Comments
 (0)