Source code for est.core.utils.pretty_print

import sys
from collections.abc import Mapping
from collections.abc import Sequence
from typing import Optional

import numpy
from pydantic import BaseModel




def tree(obj, title=None) -> str:
    """
    Pretty-print a BaseModel or dictionary as a tree.

    - Nested dicts and BaseModels are recursed.
    - NumPy arrays are summarized (shape, min, max).
    - ASCII fallback if stdout does not support UTF-8.
    """
    use_unicode = _supports_unicode()
    return _tree_recursive(obj, prefix="", use_unicode=use_unicode, title=title)



def _tree_recursive(obj, prefix="", use_unicode=True, title=None) -> str:
    pipe = "│   " if use_unicode else "|   "
    tee = "├── " if use_unicode else "|-- "
    corner = "└── " if use_unicode else "`-- "

    lines = []

    # Root node
    if prefix == "":
        if isinstance(obj, BaseModel):
            lines.append(title or type(obj).__name__)
        elif isinstance(obj, Mapping):
            lines.append(title or "root")
        else:
            lines.append(title or repr(obj))
            return "\n".join(lines)

    # Determine keys/items
    if isinstance(obj, BaseModel):
        items = [(name, getattr(obj, name)) for name in type(obj).model_fields.keys()]
    elif isinstance(obj, Mapping):
        items = list(obj.items())
    else:
        # Not a container: just print summary
        lines.append(prefix + corner + repr(obj))
        return "\n".join(lines)

    for i, (key, value) in enumerate(items):
        last = i == len(items) - 1
        connector = corner if last else tee
        line_prefix = prefix + connector

        summary = _summarize(value)
        if summary is not None:
            lines.append(f"{line_prefix}{key}: {summary}")
        elif isinstance(value, (Mapping, BaseModel)):
            lines.append(f"{line_prefix}{key}")
            extension = "    " if last else pipe
            lines.append(
                _tree_recursive(value, prefix + extension, use_unicode=use_unicode)
            )
        else:
            lines.append(f"{line_prefix}{key}: {repr(value)}")

    return "\n".join(lines)


def _summarize(value) -> Optional[str]:
    """Return a string summary for non-container types."""
    if isinstance(value, numpy.ndarray):
        if value.size > 0:
            return f"shape={value.shape} min={value.min():.3g} max={value.max():.3g}"
        else:
            return f"shape={value.shape} EMPTY"
    elif value is None:
        return "<MISSING>"
    elif isinstance(value, Sequence) and not isinstance(value, str):
        return f"len={len(value)}"
    return None


def _supports_unicode() -> bool:
    encoding = sys.stdout.encoding or ""
    return encoding.lower().startswith("utf")