Skip to content

Commit 4587f33

Browse files
authored
ENH: __repr__ for 2D DTA/TDA (#37164)
1 parent cc4a916 commit 4587f33

File tree

4 files changed

+55
-2
lines changed

4 files changed

+55
-2
lines changed

pandas/core/arrays/_mixins.py

+21
Original file line numberDiff line numberDiff line change
@@ -252,3 +252,24 @@ def _reduce(self, name: str, skipna: bool = True, **kwargs):
252252
else:
253253
msg = f"'{type(self).__name__}' does not implement reduction '{name}'"
254254
raise TypeError(msg)
255+
256+
# ------------------------------------------------------------------------
257+
258+
def __repr__(self) -> str:
259+
if self.ndim == 1:
260+
return super().__repr__()
261+
262+
from pandas.io.formats.printing import format_object_summary
263+
264+
# the short repr has no trailing newline, while the truncated
265+
# repr does. So we include a newline in our template, and strip
266+
# any trailing newlines from format_object_summary
267+
lines = [
268+
format_object_summary(x, self._formatter(), indent_for_name=False).rstrip(
269+
", \n"
270+
)
271+
for x in self
272+
]
273+
data = ",\n".join(lines)
274+
class_name = f"<{type(self).__name__}>"
275+
return f"{class_name}\n[\n{data}\n]\nShape: {self.shape}, dtype: {self.dtype}"

pandas/core/arrays/datetimes.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -566,7 +566,8 @@ def __iter__(self):
566566
tstamp : Timestamp
567567
"""
568568
if self.ndim > 1:
569-
return (self[n] for n in range(len(self)))
569+
for i in range(len(self)):
570+
yield self[i]
570571
else:
571572
# convert in chunks of 10k for efficiency
572573
data = self.asi8

pandas/core/arrays/timedeltas.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -352,7 +352,8 @@ def astype(self, dtype, copy: bool = True):
352352

353353
def __iter__(self):
354354
if self.ndim > 1:
355-
return (self[n] for n in range(len(self)))
355+
for i in range(len(self)):
356+
yield self[i]
356357
else:
357358
# convert in chunks of 10k for efficiency
358359
data = self.asi8

pandas/tests/arrays/test_datetimelike.py

+30
Original file line numberDiff line numberDiff line change
@@ -328,11 +328,41 @@ def test_iter_2d(self, arr1d):
328328
data2d = arr1d._data[:3, np.newaxis]
329329
arr2d = type(arr1d)._simple_new(data2d, dtype=arr1d.dtype)
330330
result = list(arr2d)
331+
assert len(result) == 3
331332
for x in result:
332333
assert isinstance(x, type(arr1d))
333334
assert x.ndim == 1
334335
assert x.dtype == arr1d.dtype
335336

337+
def test_repr_2d(self, arr1d):
338+
data2d = arr1d._data[:3, np.newaxis]
339+
arr2d = type(arr1d)._simple_new(data2d, dtype=arr1d.dtype)
340+
341+
result = repr(arr2d)
342+
343+
if isinstance(arr2d, TimedeltaArray):
344+
expected = (
345+
f"<{type(arr2d).__name__}>\n"
346+
"[\n"
347+
f"['{arr1d[0]._repr_base()}'],\n"
348+
f"['{arr1d[1]._repr_base()}'],\n"
349+
f"['{arr1d[2]._repr_base()}']\n"
350+
"]\n"
351+
f"Shape: (3, 1), dtype: {arr1d.dtype}"
352+
)
353+
else:
354+
expected = (
355+
f"<{type(arr2d).__name__}>\n"
356+
"[\n"
357+
f"['{arr1d[0]}'],\n"
358+
f"['{arr1d[1]}'],\n"
359+
f"['{arr1d[2]}']\n"
360+
"]\n"
361+
f"Shape: (3, 1), dtype: {arr1d.dtype}"
362+
)
363+
364+
assert result == expected
365+
336366
def test_setitem(self):
337367
data = np.arange(10, dtype="i8") * 24 * 3600 * 10 ** 9
338368
arr = self.array_cls(data, freq="D")

0 commit comments

Comments
 (0)