diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py index b508ded18..33aa35a22 100644 --- a/narwhals/dataframe.py +++ b/narwhals/dataframe.py @@ -19,6 +19,7 @@ from narwhals.translate import to_native from narwhals.utils import find_stacklevel from narwhals.utils import flatten +from narwhals.utils import generate_repr from narwhals.utils import is_sequence_but_not_str from narwhals.utils import parse_version @@ -414,18 +415,7 @@ def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray: return self._compliant_frame.__array__(dtype, copy=copy) def __repr__(self) -> str: # pragma: no cover - header = " Narwhals DataFrame " - length = len(header) - return ( - "┌" - + "─" * length - + "┐\n" - + f"|{header}|\n" - + "| Use `.to_native` to see native output |\n" - + "└" - + "─" * length - + "┘" - ) + return generate_repr("Narwhals DataFrame", self.to_native().__repr__()) def __arrow_c_stream__(self, requested_schema: object | None = None) -> object: """Export a DataFrame via the Arrow PyCapsule Interface. @@ -3581,18 +3571,7 @@ def __init__( raise AssertionError(msg) def __repr__(self) -> str: # pragma: no cover - header = " Narwhals LazyFrame " - length = len(header) - return ( - "┌" - + "─" * length - + "┐\n" - + f"|{header}|\n" - + "| Use `.to_native` to see native output |\n" - + "└" - + "─" * length - + "┘" - ) + return generate_repr("Narwhals LazyFrame", self.to_native().__repr__()) @property def implementation(self) -> Implementation: @@ -3640,11 +3619,12 @@ def collect(self) -> DataFrame[Any]: ... } ... ) >>> lf = nw.from_native(lf_pl) - >>> lf - ┌───────────────────────────────────────┐ - | Narwhals LazyFrame | - | Use `.to_native` to see native output | - └───────────────────────────────────────┘ + >>> lf # doctest:+ELLIPSIS + ┌─────────────────────────────┐ + | Narwhals LazyFrame | + |-----------------------------| + |>> df = lf.group_by("a").agg(nw.all().sum()).collect() >>> df.to_native().sort("a") shape: (3, 3) diff --git a/narwhals/series.py b/narwhals/series.py index 8f15ff0ce..de0e64396 100644 --- a/narwhals/series.py +++ b/narwhals/series.py @@ -15,6 +15,7 @@ from narwhals.dtypes import _validate_dtype from narwhals.typing import IntoSeriesT from narwhals.utils import _validate_rolling_arguments +from narwhals.utils import generate_repr from narwhals.utils import parse_version if TYPE_CHECKING: @@ -404,18 +405,7 @@ def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Se return function(self, *args, **kwargs) def __repr__(self) -> str: # pragma: no cover - header = " Narwhals Series " - length = len(header) - return ( - "┌" - + "─" * length - + "┐\n" - + f"|{header}|\n" - + "| Use `.to_native()` to see native output |\n" - + "└" - + "─" * length - + "┘" - ) + return generate_repr("Narwhals Series", self.to_native().__repr__()) def __len__(self) -> int: return len(self._compliant_series) diff --git a/narwhals/stable/v1/__init__.py b/narwhals/stable/v1/__init__.py index f7705713f..22afc687d 100644 --- a/narwhals/stable/v1/__init__.py +++ b/narwhals/stable/v1/__init__.py @@ -430,11 +430,12 @@ def collect(self) -> DataFrame[Any]: ... } ... ) >>> lf = nw.from_native(lf_pl) - >>> lf - ┌───────────────────────────────────────┐ - | Narwhals LazyFrame | - | Use `.to_native` to see native output | - └───────────────────────────────────────┘ + >>> lf # doctest:+ELLIPSIS + ┌─────────────────────────────┐ + | Narwhals LazyFrame | + |-----------------------------| + |>> df = lf.group_by("a").agg(nw.all().sum()).collect() >>> df.to_native().sort("a") shape: (3, 3) diff --git a/narwhals/utils.py b/narwhals/utils.py index b6337cb8e..2125d46c4 100644 --- a/narwhals/utils.py +++ b/narwhals/utils.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os import re from enum import Enum from enum import auto @@ -960,3 +961,35 @@ def _validate_rolling_arguments( min_periods = window_size return window_size, min_periods + + +def generate_repr(header: str, native_repr: str) -> str: + try: + terminal_width = os.get_terminal_size().columns + except OSError: + terminal_width = 80 + native_lines = native_repr.splitlines() + max_native_width = max(len(line) for line in native_lines) + + if max_native_width + 2 < terminal_width: + length = max(max_native_width, len(header)) + output = f"┌{'─'*length}┐\n" + header_extra = length - len(header) + output += ( + f"|{' '*(header_extra//2)}{header}{' '*(header_extra//2 + header_extra%2)}|\n" + ) + output += f"|{'-'*(length)}|\n" + start_extra = (length - max_native_width) // 2 + end_extra = (length - max_native_width) // 2 + (length - max_native_width) % 2 + for line in native_lines: + output += f"|{' '*(start_extra)}{line}{' '*(end_extra + max_native_width - len(line))}|\n" + output += f"└{'─' * length}┘" + return output + + diff = 39 - len(header) + return ( + f"┌{'─' * (39)}┐\n" + f"|{' '*(diff//2)}{header}{' '*(diff//2+diff%2)}|\n" + "| Use `.to_native` to see native output |\n└" + f"{'─' * 39}┘" + ) diff --git a/tests/repr_test.py b/tests/repr_test.py new file mode 100644 index 000000000..40cd51dca --- /dev/null +++ b/tests/repr_test.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import pandas as pd +import pytest + +import narwhals.stable.v1 as nw + + +def test_repr() -> None: + duckdb = pytest.importorskip("duckdb") + df = pd.DataFrame({"a": [1, 2, 3], "b": ["fdaf", "fda", "cf"]}) + result = nw.from_native(df).__repr__() + expected = ( + "┌──────────────────┐\n" + "|Narwhals DataFrame|\n" + "|------------------|\n" + "| a b |\n" + "| 0 1 fdaf |\n" + "| 1 2 fda |\n" + "| 2 3 cf |\n" + "└──────────────────┘" + ) + assert result == expected + result = nw.from_native(df).lazy().__repr__() + expected = ( + "┌──────────────────┐\n" + "|Narwhals LazyFrame|\n" + "|------------------|\n" + "| a b |\n" + "| 0 1 fdaf |\n" + "| 1 2 fda |\n" + "| 2 3 cf |\n" + "└──────────────────┘" + ) + assert result == expected + result = nw.from_native(df)["a"].__repr__() + expected = ( + "┌─────────────────────┐\n" + "| Narwhals Series |\n" + "|---------------------|\n" + "|0 1 |\n" + "|1 2 |\n" + "|2 3 |\n" + "|Name: a, dtype: int64|\n" + "└─────────────────────┘" + ) + assert result == expected + result = nw.from_native(duckdb.table("df")).__repr__() + expected = ( + "┌───────────────────┐\n" + "|Narwhals DataFrame |\n" + "|-------------------|\n" + "|┌───────┬─────────┐|\n" + "|│ a │ b │|\n" + "|│ int64 │ varchar │|\n" + "|├───────┼─────────┤|\n" + "|│ 1 │ fdaf │|\n" + "|│ 2 │ fda │|\n" + "|│ 3 │ cf │|\n" + "|└───────┴─────────┘|\n" + "└───────────────────┘" + ) + assert result == expected + # Make something wider than the terminal size + df = pd.DataFrame({"a": [1, 2, 3], "b": ["fdaf" * 100, "fda", "cf"]}) + result = nw.from_native(duckdb.table("df")).__repr__() + expected = ( + "┌───────────────────────────────────────┐\n" + "| Narwhals DataFrame |\n" + "| Use `.to_native` to see native output |\n" + "└───────────────────────────────────────┘" + ) + assert result == expected