Add rich repr (#409)

* Add rich repr

Closes #393

* Fixes.

* Cleanup + add tests

* Try again

* Fix tests

* revert unnecessary change

* Right justify

* de-emphasize empty rows; highlight CF names more

* Refactor

* Small cleanup

* Small tweaks + solarized colors

* fix  tests.

* fix typing

* Fix test

* Add back type ignores.

Revert "fix typing"

This reverts commit 0f9c57e.

* Update docs

* Add as optional dep

* Add image.

* WIP conventions

* Merge upstream/main

* Add whats-new note

Author: dcherian

cf_xarray/accessor.py

@@ -36,6 +36,14 @@
     grid_mapping_var_criteria,
     regex,
 )
+from .formatting import (
+    _format_coordinates,
+    # _format_conventions,
+    _format_data_vars,
+    _format_flags,
+    _format_roles,
+    _maybe_panel,
+)
 from .helpers import _guess_bounds_1d, _guess_bounds_2d, bounds_to_vertices
 from .options import OPTIONS
 from .utils import (
@@ -602,6 +610,11 @@ def _getattr(
         An extra decorator, if necessary. This is used by _CFPlotMethods to set default
         kwargs based on CF attributes.
     """
+
+    # UGH. this seems unavoidable because I'm overriding getattr
+    if attr in ["_repr_html_", "__rich__", "__rich_console__"]:
+        raise AttributeError
+
     try:
         attribute: Mapping | Callable = getattr(obj, attr)
     except AttributeError:
@@ -669,7 +682,9 @@ def wrapper(*args, **kwargs):
 
         return result
 
-    wrapper.__doc__ = _build_docstring(func) + wrapper.__doc__
+    # handle rich
+    if wrapper.__doc__:
+        wrapper.__doc__ = _build_docstring(func) + wrapper.__doc__
 
     return wrapper
 
@@ -1399,90 +1414,45 @@ def describe(self):
         print(repr(self))
 
     def __repr__(self):
+        return ("".join(self._generate_repr(rich=False))).rstrip()
 
-        coords = self._obj.coords
-        dims = self._obj.dims
-
-        def make_text_section(subtitle, attr, valid_values=None, default_keys=None):
-
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore")
-                try:
-                    vardict = getattr(self, attr, {})
-                except ValueError:
-                    vardict = {}
-
-            star = " * "
-            tab = len(star) * " "
-            subtitle = f"- {subtitle}:"
-
-            # Sort keys if there aren't extra keys,
-            # preserve default keys order otherwise.
-            default_keys = [] if not default_keys else list(default_keys)
-            extra_keys = list(set(vardict) - set(default_keys))
-            ordered_keys = sorted(vardict) if extra_keys else default_keys
-            vardict = {key: vardict[key] for key in ordered_keys if key in vardict}
-
-            # Keep only valid values (e.g., coords or data_vars)
-            if valid_values is not None:
-                vardict = {
-                    key: set(value).intersection(valid_values)
-                    for key, value in vardict.items()
-                    if set(value).intersection(valid_values)
-                }
+    def __rich__(self):
+        from rich.console import Group
 
-            # Star for keys with dims only, tab otherwise
-            rows = [
-                f"{star if set(value) <= set(dims) else tab}{key}: {sorted(value)}"
-                for key, value in vardict.items()
-            ]
+        return Group(*self._generate_repr(rich=True))
 
-            # Append missing default keys followed by n/a
-            if default_keys:
-                missing_keys = [key for key in default_keys if key not in vardict]
-                if missing_keys:
-                    rows += [tab + ", ".join(missing_keys) + ": n/a"]
-            elif not rows:
-                rows = [tab + "n/a"]
-
-            # Add subtitle to the first row, align other rows
-            rows = [
-                "\n" + subtitle + row if i == 0 else len(subtitle) * " " + row
-                for i, row in enumerate(rows)
-            ]
+    def _generate_repr(self, rich=False):
+        dims = self._obj.dims
+        coords = self._obj.coords
 
-            return "\n".join(rows) + "\n"
+        # if self._obj._attrs:
+        #     conventions = self._obj.attrs.pop("Conventions", None)
+        #     if conventions:
+        #         yield _format_conventions(conventions, rich)
 
         if isinstance(self._obj, DataArray) and self._obj.cf.is_flag_variable:
-            flag_dict = create_flag_dict(self._obj)
-            text = f"CF Flag variable with mapping:\n\t{flag_dict!r}\n\n"
-        else:
-            text = ""
+            yield _maybe_panel(
+                _format_flags(self, rich), title="Flag Variable", rich=rich
+            )
 
         if self.cf_roles:
-            text += make_text_section("CF Roles", "cf_roles")
-            text += "\n"
-
-        text += "Coordinates:"
-        text += make_text_section("CF Axes", "axes", coords, _AXIS_NAMES)
-        text += make_text_section("CF Coordinates", "coordinates", coords, _COORD_NAMES)
-        text += make_text_section(
-            "Cell Measures", "cell_measures", coords, _CELL_MEASURES
+            yield _maybe_panel(
+                _format_roles(self, dims, rich),
+                title="Discrete Sampling Geometry",
+                rich=rich,
+            )
+
+        yield _maybe_panel(
+            _format_coordinates(self, dims, coords, rich),
+            title="Coordinates",
+            rich=rich,
         )
-        text += make_text_section("Standard Names", "standard_names", coords)
-        text += make_text_section("Bounds", "bounds", coords)
         if isinstance(self._obj, Dataset):
-            text += make_text_section("Grid Mappings", "grid_mapping_names", coords)
-            data_vars = self._obj.data_vars
-            text += "\nData Variables:"
-            text += make_text_section(
-                "Cell Measures", "cell_measures", data_vars, _CELL_MEASURES
+            yield _maybe_panel(
+                _format_data_vars(self, self._obj.data_vars, rich),
+                title="Data Variables",
+                rich=rich,
             )
-            text += make_text_section("Standard Names", "standard_names", data_vars)
-            text += make_text_section("Bounds", "bounds", data_vars)
-            text += make_text_section("Grid Mappings", "grid_mapping_names", data_vars)
-
-        return text
 
     def keys(self) -> set[Hashable]:
         """

cf_xarray/formatting.py

@@ -0,0 +1,194 @@
+import warnings
+from typing import Dict, Hashable, Iterable, List
+
+STAR = " * "
+TAB = len(STAR) * " "
+
+
+def _format_missing_row(row: str, rich: bool) -> str:
+    if rich:
+        return f"[grey62]{row}[/grey62]"
+    else:
+        return row
+
+
+def _format_varname(name, rich: bool):
+    return name
+
+
+def _format_subtitle(name: str, rich: bool) -> str:
+    if rich:
+        return f"[bold]{name}[/bold]"
+    else:
+        return name
+
+
+def _format_cf_name(name: str, rich: bool) -> str:
+    if rich:
+        return f"[color(33)]{name}[/color(33)]"
+    else:
+        return name
+
+
+def make_text_section(
+    accessor,
+    subtitle: str,
+    attr: str,
+    dims=None,
+    valid_values=None,
+    default_keys=None,
+    rich: bool = False,
+):
+
+    from .accessor import sort_maybe_hashable
+
+    if dims is None:
+        dims = []
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        try:
+            vardict: Dict[str, Iterable[Hashable]] = getattr(accessor, attr, {})
+        except ValueError:
+            vardict = {}
+
+    # Sort keys if there aren't extra keys,
+    # preserve default keys order otherwise.
+    default_keys = [] if not default_keys else list(default_keys)
+    extra_keys = list(set(vardict) - set(default_keys))
+    ordered_keys = sorted(vardict) if extra_keys else default_keys
+    vardict = {key: vardict[key] for key in ordered_keys if key in vardict}
+
+    # Keep only valid values (e.g., coords or data_vars)
+    if valid_values is not None:
+        vardict = {
+            key: set(value).intersection(valid_values)
+            for key, value in vardict.items()
+            if set(value).intersection(valid_values)
+        }
+
+    # Star for keys with dims only, tab otherwise
+    rows = [
+        (
+            f"{STAR if dims and set(value) <= set(dims) else TAB}"
+            f"{_format_cf_name(key, rich)}: "
+            f"{_format_varname(sort_maybe_hashable(value), rich)}"
+        )
+        for key, value in vardict.items()
+    ]
+
+    # Append missing default keys followed by n/a
+    if default_keys:
+        missing_keys = [key for key in default_keys if key not in vardict]
+        if missing_keys:
+            rows.append(
+                _format_missing_row(TAB + ", ".join(missing_keys) + ": n/a", rich)
+            )
+    elif not rows:
+        rows.append(_format_missing_row(TAB + "n/a", rich))
+
+    return _print_rows(subtitle, rows, rich)
+
+
+def _print_rows(subtitle: str, rows: List[str], rich: bool):
+    subtitle = f"{subtitle.rjust(20)}:"
+
+    # Add subtitle to the first row, align other rows
+    rows = [
+        _format_subtitle(subtitle, rich=rich) + row
+        if i == 0
+        else len(subtitle) * " " + row
+        for i, row in enumerate(rows)
+    ]
+
+    return "\n".join(rows) + "\n\n"
+
+
+def _format_conventions(string: str, rich: bool):
+    row = _print_rows(
+        subtitle="Conventions",
+        rows=[_format_cf_name(TAB + string, rich=rich)],
+        rich=rich,
+    )
+    if rich:
+        row = row.rstrip()
+    return row
+
+
+def _maybe_panel(textgen, title: str, rich: bool):
+    text = "".join(textgen)
+    if rich:
+        from rich.panel import Panel
+
+        return Panel(
+            f"[color(241)]{text.rstrip()}[/color(241)]",
+            expand=True,
+            title_align="left",
+            title=f"[bold][color(244)]{title}[/bold][/color(244)]",
+            highlight=True,
+            width=100,
+        )
+    else:
+        return title + ":\n" + text
+
+
+def _format_flags(accessor, rich):
+    from .accessor import create_flag_dict
+
+    flag_dict = create_flag_dict(accessor._obj)
+    rows = [
+        f"{TAB}{_format_varname(v, rich)}: {_format_cf_name(k, rich)}"
+        for k, v in flag_dict.items()
+    ]
+    return _print_rows("Flag Meanings", rows, rich)
+
+
+def _format_roles(accessor, dims, rich):
+    yield make_text_section(accessor, "CF Roles", "cf_roles", dims=dims, rich=rich)
+
+
+def _format_coordinates(accessor, dims, coords, rich):
+    from .accessor import _AXIS_NAMES, _CELL_MEASURES, _COORD_NAMES
+
+    yield make_text_section(
+        accessor, "CF Axes", "axes", dims, coords, _AXIS_NAMES, rich=rich
+    )
+    yield make_text_section(
+        accessor, "CF Coordinates", "coordinates", dims, coords, _COORD_NAMES, rich=rich
+    )
+    yield make_text_section(
+        accessor,
+        "Cell Measures",
+        "cell_measures",
+        dims,
+        coords,
+        _CELL_MEASURES,
+        rich=rich,
+    )
+    yield make_text_section(
+        accessor, "Standard Names", "standard_names", dims, coords, rich=rich
+    )
+    yield make_text_section(accessor, "Bounds", "bounds", dims, coords, rich=rich)
+    yield make_text_section(
+        accessor, "Grid Mappings", "grid_mapping_names", dims, coords, rich=rich
+    )
+
+
+def _format_data_vars(accessor, data_vars, rich):
+    from .accessor import _CELL_MEASURES
+
+    yield make_text_section(
+        accessor,
+        "Cell Measures",
+        "cell_measures",
+        None,
+        data_vars,
+        _CELL_MEASURES,
+        rich=rich,
+    )
+    yield make_text_section(
+        accessor, "Standard Names", "standard_names", None, data_vars, rich=rich
+    )
+    yield make_text_section(accessor, "Bounds", "bounds", None, data_vars, rich=rich)
+    yield make_text_section(
+        accessor, "Grid Mappings", "grid_mapping_names", None, data_vars, rich=rich
+    )

cf_xarray/tests/__init__.py

@@ -67,4 +67,5 @@ def LooseVersion(vstring):
 has_scipy, requires_scipy = _importorskip("scipy")
 has_shapely, requires_shapely = _importorskip("shapely")
 has_pint, requires_pint = _importorskip("pint")
+_, requires_rich = _importorskip("rich")
 has_regex, requires_regex = _importorskip("regex")