Skip to content

Commit

Permalink
ENH: __repr__ for 2D DTA/TDA (#37164)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored Nov 2, 2020
1 parent cc4a916 commit 4587f33
Show file tree
Hide file tree
Showing 4 changed files with 55 additions and 2 deletions.
21 changes: 21 additions & 0 deletions pandas/core/arrays/_mixins.py
Original file line number Diff line number Diff line change
Expand Up @@ -252,3 +252,24 @@ def _reduce(self, name: str, skipna: bool = True, **kwargs):
else:
msg = f"'{type(self).__name__}' does not implement reduction '{name}'"
raise TypeError(msg)

# ------------------------------------------------------------------------

def __repr__(self) -> str:
if self.ndim == 1:
return super().__repr__()

from pandas.io.formats.printing import format_object_summary

# the short repr has no trailing newline, while the truncated
# repr does. So we include a newline in our template, and strip
# any trailing newlines from format_object_summary
lines = [
format_object_summary(x, self._formatter(), indent_for_name=False).rstrip(
", \n"
)
for x in self
]
data = ",\n".join(lines)
class_name = f"<{type(self).__name__}>"
return f"{class_name}\n[\n{data}\n]\nShape: {self.shape}, dtype: {self.dtype}"
3 changes: 2 additions & 1 deletion pandas/core/arrays/datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -566,7 +566,8 @@ def __iter__(self):
tstamp : Timestamp
"""
if self.ndim > 1:
return (self[n] for n in range(len(self)))
for i in range(len(self)):
yield self[i]
else:
# convert in chunks of 10k for efficiency
data = self.asi8
Expand Down
3 changes: 2 additions & 1 deletion pandas/core/arrays/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -352,7 +352,8 @@ def astype(self, dtype, copy: bool = True):

def __iter__(self):
if self.ndim > 1:
return (self[n] for n in range(len(self)))
for i in range(len(self)):
yield self[i]
else:
# convert in chunks of 10k for efficiency
data = self.asi8
Expand Down
30 changes: 30 additions & 0 deletions pandas/tests/arrays/test_datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -328,11 +328,41 @@ def test_iter_2d(self, arr1d):
data2d = arr1d._data[:3, np.newaxis]
arr2d = type(arr1d)._simple_new(data2d, dtype=arr1d.dtype)
result = list(arr2d)
assert len(result) == 3
for x in result:
assert isinstance(x, type(arr1d))
assert x.ndim == 1
assert x.dtype == arr1d.dtype

def test_repr_2d(self, arr1d):
data2d = arr1d._data[:3, np.newaxis]
arr2d = type(arr1d)._simple_new(data2d, dtype=arr1d.dtype)

result = repr(arr2d)

if isinstance(arr2d, TimedeltaArray):
expected = (
f"<{type(arr2d).__name__}>\n"
"[\n"
f"['{arr1d[0]._repr_base()}'],\n"
f"['{arr1d[1]._repr_base()}'],\n"
f"['{arr1d[2]._repr_base()}']\n"
"]\n"
f"Shape: (3, 1), dtype: {arr1d.dtype}"
)
else:
expected = (
f"<{type(arr2d).__name__}>\n"
"[\n"
f"['{arr1d[0]}'],\n"
f"['{arr1d[1]}'],\n"
f"['{arr1d[2]}']\n"
"]\n"
f"Shape: (3, 1), dtype: {arr1d.dtype}"
)

assert result == expected

def test_setitem(self):
data = np.arange(10, dtype="i8") * 24 * 3600 * 10 ** 9
arr = self.array_cls(data, freq="D")
Expand Down

0 comments on commit 4587f33

Please sign in to comment.