ezpz.log

Logging utilities built on top of Rich.

ezpz/log/init.py

Console

Bases: Console

Extends rich Console class.

Source code in src/ezpz/log/console.py

class Console(rich_console.Console):
    """Extends rich Console class."""

    def __init__(self, *args: str, redirect: bool = True, **kwargs: Any) -> None:
        """
        enrich console does soft-wrapping by default and this diverge from
        original rich console which does not, creating hard-wraps instead.
        """
        self.redirect = redirect

        if "soft_wrap" not in kwargs:
            kwargs["soft_wrap"] = True

        if "theme" not in kwargs:
            kwargs["theme"] = get_theme()

        if "markup" not in kwargs:
            kwargs["markup"] = True

        if "width" not in kwargs:
            kwargs["width"] = 55510

        # Unless user already mentioning terminal preference, we use our
        # heuristic to make an informed decision.
        if "force_terminal" not in kwargs:
            kwargs["force_terminal"] = should_do_markup(
                stream=kwargs.get("file", sys.stdout)
            )
        if "force_jupyter" not in kwargs:
            kwargs["force_jupyter"] = is_interactive()

        super().__init__(*args, **kwargs)
        self.extended = True

        if self.redirect:
            if not hasattr(sys.stdout, "rich_proxied_file"):
                sys.stdout = FileProxy(self, sys.stdout)  # type: ignore
            if not hasattr(sys.stderr, "rich_proxied_file"):
                sys.stderr = FileProxy(self, sys.stderr)  # type: ignore

    # https://github.com/python/mypy/issues/4441
    def print(self, *args, **kwargs) -> None:  # type: ignore
        """Print override that respects user soft_wrap preference."""
        # Currently rich is unable to render ANSI escapes with print so if
        # we detect their presence, we decode them.
        # https://github.com/willmcgugan/rich/discussions/404
        if args and isinstance(args[0], str) and "\033" in args[0]:
            text = format(*args) + "\n"
            decoder = AnsiDecoder()
            args = list(decoder.decode(text))  # type: ignore
        super().print(*args, **kwargs)

__init__(*args, redirect=True, **kwargs)

enrich console does soft-wrapping by default and this diverge from original rich console which does not, creating hard-wraps instead.

Source code in src/ezpz/log/console.py

def __init__(self, *args: str, redirect: bool = True, **kwargs: Any) -> None:
    """
    enrich console does soft-wrapping by default and this diverge from
    original rich console which does not, creating hard-wraps instead.
    """
    self.redirect = redirect

    if "soft_wrap" not in kwargs:
        kwargs["soft_wrap"] = True

    if "theme" not in kwargs:
        kwargs["theme"] = get_theme()

    if "markup" not in kwargs:
        kwargs["markup"] = True

    if "width" not in kwargs:
        kwargs["width"] = 55510

    # Unless user already mentioning terminal preference, we use our
    # heuristic to make an informed decision.
    if "force_terminal" not in kwargs:
        kwargs["force_terminal"] = should_do_markup(
            stream=kwargs.get("file", sys.stdout)
        )
    if "force_jupyter" not in kwargs:
        kwargs["force_jupyter"] = is_interactive()

    super().__init__(*args, **kwargs)
    self.extended = True

    if self.redirect:
        if not hasattr(sys.stdout, "rich_proxied_file"):
            sys.stdout = FileProxy(self, sys.stdout)  # type: ignore
        if not hasattr(sys.stderr, "rich_proxied_file"):
            sys.stderr = FileProxy(self, sys.stderr)  # type: ignore

print(*args, **kwargs)

Print override that respects user soft_wrap preference.

Source code in src/ezpz/log/console.py

def print(self, *args, **kwargs) -> None:  # type: ignore
    """Print override that respects user soft_wrap preference."""
    # Currently rich is unable to render ANSI escapes with print so if
    # we detect their presence, we decode them.
    # https://github.com/willmcgugan/rich/discussions/404
    if args and isinstance(args[0], str) and "\033" in args[0]:
        text = format(*args) + "\n"
        decoder = AnsiDecoder()
        args = list(decoder.decode(text))  # type: ignore
    super().print(*args, **kwargs)

FluidLogRender

Renders log by not using columns and avoiding any wrapping.

Source code in src/ezpz/log/handler.py

class FluidLogRender:  # pylint: disable=too-few-public-methods
    """Renders log by not using columns and avoiding any wrapping."""

    def __init__(
        self,
        show_time: bool = True,
        show_level: bool = True,
        show_path: bool = True,
        time_format: str = "%Y-%m-%d %H:%M:%S.%f",
        link_path: Optional[bool] = False,
    ) -> None:
        self.show_time = show_time
        self.show_level = show_level
        self.show_path = show_path
        self.time_format = time_format
        self.link_path = link_path
        self._last_time: Optional[str] = None
        self.colorized = use_colored_logs()
        self.styles = get_styles()

    def __call__(  # pylint: disable=too-many-arguments
        self,
        console: Console,  # type: ignore
        renderables: Iterable[ConsoleRenderable],
        log_time: Optional[datetime] = None,
        time_format: str = "%Y-%m-%d %H:%M:%S.%f",
        level: TextType = "",
        path: Optional[str] = None,
        line_no: Optional[int] = None,
        link_path: Optional[str] = None,
        funcName: Optional[str] = None,
    ) -> Text:
        result = Text()
        if self.show_time:
            log_time = datetime.now() if log_time is None else log_time
            log_time_display = log_time.strftime(time_format or self.time_format)
            d, t = log_time_display.split(" ")
            result += Text("[", style=self.styles.get("log.brace", ""))
            result += Text(f"{d} ", style=self.styles.get("logging.date", ""))
            result += Text(t, style=self.styles.get("logging.time", ""))
            result += Text("]", style=self.styles.get("log.brace", ""))
            # result += Text(log_time_display, style=self.styles['logging.time'])
            self._last_time = log_time_display
        if self.show_level:
            if isinstance(level, Text):
                lstr = level.plain.rstrip(" ")[0]
                if self.colorized:
                    style = level.spans[0].style
                else:
                    style = Style.null()
                level.spans = [Span(0, len(lstr), style)]
                ltext = Text("[", style=self.styles.get("log.brace", ""))
                ltext.append(Text(f"{lstr}", style=style))
                ltext.append(Text("]", style=self.styles.get("log.brace", "")))
                # ltext = Text(f'[{lstr}]', style=style)
            elif isinstance(level, str):
                lstr = level.rstrip(" ")[0]
                style = f"logging.level.{str(lstr)}" if self.colorized else Style.null()
                ltext = Text("[", style=self.styles.get("log.brace", ""))
                ltext = Text(f"{lstr}", style=style)  # f"logging.level.{str(lstr)}")
                ltext.append(Text("]", style=self.styles.get("log.brace", "")))
            result += ltext
        if self.show_path and path:
            path_text = Text("[", style=self.styles.get("log.brace", ""))
            text_arr = []
            # try:
            # fp = Path(path)
            # parent = fp.parent.as_posix().split("/")[-1]
            # remainder = fp.stem
            parent, remainder = path.split("/")
            # except Exception:
            #     import ezpz
            #     ezpz.breakpoint(0)
            if "." in remainder:
                module, *fn = remainder.split(".")
                fn = ".".join(fn)
            else:
                module = remainder
                fn = None
            if funcName is not None:
                fn = funcName
            text_arr += [
                Text(
                    f"{parent}", style="log.parent"
                ),  # self.styles.get('log.pa', '')),
                Text("/"),
                Text(f"{module}", style="log.path"),
            ]
            if line_no:
                text_arr += [
                    Text(":", style=self.styles.get("log.colon", "")),
                    Text(
                        f"{line_no}",
                        style=self.styles.get("log.linenumber", ""),
                    ),
                ]
                # text_arr.append(Text(':', style=self.styles.get('log.colon', '')))
                # text_arr.append(
                #     Text(f'{line_no}', style=self.styles.get('log.linenumber', '')),
                # )
            if fn is not None:
                text_arr += [
                    Text(":", style="log.colon"),
                    Text(
                        f"{fn}",
                        style="repr.function",  # self.styles.get('repr.inspect.def', 'json.key'),
                    ),
                ]
            path_text.append(Text.join(Text(""), text_arr))

            # for t in text_arr:
            #     path_text.append(t)
            # path_text.append(
            #     Text(f'{parent}', style='cyan'),
            # )
            # path_text.append(Text('/'))
            # path_text.append(
            #     remainder,
            #     style=STYLES.get('log.path', ''),
            # )
            # if fn is not None:
            #     path_text += [Text('.'), Text(f'{fn}', style='cyan')]

            path_text.append("]", style=self.styles.get("log.brace", ""))
            result += path_text
        result += Text(" ", style=self.styles.get("repr.dash", ""))
        for elem in renderables:
            if ANSI_ESCAPE_PATTERN.search(str(elem)):
                # If the element is already ANSI formatted, append it directly
                result += Text.from_ansi(str(elem))
            else:
                result += elem
        return result

RichHandler

Bases: RichHandler

Enriched handler that does not wrap.

Source code in src/ezpz/log/handler.py

class RichHandler(OriginalRichHandler):
    """Enriched handler that does not wrap."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        if "console" not in kwargs:
            console = get_console(
                redirect=False,
                width=9999,
                markup=use_colored_logs(),
                soft_wrap=False,
            )
            kwargs["console"] = console
            self.__console = console
        else:
            self.__console = kwargs["console"]
        super().__init__(*args, **kwargs)
        # RichHandler constructor does not allow custom renderer
        # https://github.com/willmcgugan/rich/issues/438
        self._log_render = FluidLogRender(
            show_time=kwargs.get("show_time", True),
            show_level=kwargs.get("show_level", True),
            show_path=kwargs.get("show_path", True),
        )  # type: ignore

    def render(
        self,
        *,
        record: LogRecord,
        traceback: Optional[Any],
        message_renderable: "ConsoleRenderable",
    ) -> "ConsoleRenderable":
        """Render log for display.

        Args:
            record (LogRecord): logging Record.
            traceback (Optional[Traceback]): Traceback instance or None for no Traceback.
            message_renderable (ConsoleRenderable): Renderable (typically Text) containing log message contents.

        Returns:
            ConsoleRenderable: Renderable to display log.
        """
        fp = getattr(record, "pathname", None)
        parent = Path(fp).parent.as_posix().split("/")[-1] if fp else None
        module = getattr(record, "module", None)
        name = getattr(record, "name", None)
        funcName = getattr(record, "funcName", None)
        parr = []
        if fp is not None:
            fp = Path(fp)
            parent = fp.parent.as_posix().split("/")[-1]
            parr.append(parent)
        if module is not None:
            parr.append(module)

        if name is not None and parent is not None and f"{parent}.{module}" != name:
            parr.append(name)
        pstr = "/".join([parr[0], ".".join(parr[1:])])
        level = self.get_level_text(record)
        time_format = None if self.formatter is None else self.formatter.datefmt
        default_time_fmt = "%Y-%m-%d %H:%M:%S,%f"
        # default_time_fmt = "%Y-%m-%d %H:%M:%S"  # .%f'
        time_format = time_format if time_format else default_time_fmt
        log_time = datetime.fromtimestamp(record.created)

        log_renderable = self._log_render(
            self.__console,
            (
                [message_renderable]
                if not traceback
                else [message_renderable, traceback]
            ),
            log_time=log_time,
            time_format=time_format,
            level=level,
            path=pstr,  # getattr(record, "pathname", None),
            line_no=record.lineno,
            link_path=record.pathname if self.enable_link_path else None,
            funcName=record.funcName,
        )
        return log_renderable

render(*, record, traceback, message_renderable)

Render log for display.

Parameters:

Name	Type	Description	Default
record	LogRecord	logging Record.	required
traceback	Optional[Traceback]	Traceback instance or None for no Traceback.	required
message_renderable	ConsoleRenderable	Renderable (typically Text) containing log message contents.	required

Returns:

Name	Type	Description
ConsoleRenderable	ConsoleRenderable	Renderable to display log.

Source code in src/ezpz/log/handler.py

def render(
    self,
    *,
    record: LogRecord,
    traceback: Optional[Any],
    message_renderable: "ConsoleRenderable",
) -> "ConsoleRenderable":
    """Render log for display.

    Args:
        record (LogRecord): logging Record.
        traceback (Optional[Traceback]): Traceback instance or None for no Traceback.
        message_renderable (ConsoleRenderable): Renderable (typically Text) containing log message contents.

    Returns:
        ConsoleRenderable: Renderable to display log.
    """
    fp = getattr(record, "pathname", None)
    parent = Path(fp).parent.as_posix().split("/")[-1] if fp else None
    module = getattr(record, "module", None)
    name = getattr(record, "name", None)
    funcName = getattr(record, "funcName", None)
    parr = []
    if fp is not None:
        fp = Path(fp)
        parent = fp.parent.as_posix().split("/")[-1]
        parr.append(parent)
    if module is not None:
        parr.append(module)

    if name is not None and parent is not None and f"{parent}.{module}" != name:
        parr.append(name)
    pstr = "/".join([parr[0], ".".join(parr[1:])])
    level = self.get_level_text(record)
    time_format = None if self.formatter is None else self.formatter.datefmt
    default_time_fmt = "%Y-%m-%d %H:%M:%S,%f"
    # default_time_fmt = "%Y-%m-%d %H:%M:%S"  # .%f'
    time_format = time_format if time_format else default_time_fmt
    log_time = datetime.fromtimestamp(record.created)

    log_renderable = self._log_render(
        self.__console,
        (
            [message_renderable]
            if not traceback
            else [message_renderable, traceback]
        ),
        log_time=log_time,
        time_format=time_format,
        level=level,
        path=pstr,  # getattr(record, "pathname", None),
        line_no=record.lineno,
        link_path=record.pathname if self.enable_link_path else None,
        funcName=record.funcName,
    )
    return log_renderable

get_active_enrich_handlers(logger)

Return (index, handler) pairs for active RichHandler instances.

Source code in src/ezpz/log/__init__.py

def get_active_enrich_handlers(logger: logging.Logger) -> list:
    """Return ``(index, handler)`` pairs for active ``RichHandler`` instances."""
    from ezpz.log.handler import RichHandler as EnrichHandler

    return [
        (idx, h)
        for idx, h in enumerate(logger.handlers)
        if isinstance(h, EnrichHandler)
    ]

get_console_from_logger(logger)

Return the Console attached to logger or synthesise a new one.

Source code in src/ezpz/log/__init__.py

def get_console_from_logger(logger: logging.Logger) -> Console:
    """Return the ``Console`` attached to *logger* or synthesise a new one."""
    from ezpz.log.handler import RichHandler as EnrichHandler

    for handler in logger.handlers:
        if isinstance(handler, (RichHandler, EnrichHandler)):
            return handler.console  # type: ignore
    from ezpz.log.console import get_console

    return get_console()

get_enrich_logging_config_as_yaml(name='enrich', level='INFO')

Render the Enrich logging YAML snippet with the requested name/level.

Source code in src/ezpz/log/__init__.py

def get_enrich_logging_config_as_yaml(name: str = "enrich", level: str = "INFO") -> str:
    """Render the Enrich logging YAML snippet with the requested name/level."""
    return rf"""
    ---
    # version: 1
    handlers:
      {name}:
        (): ezpz.log.handler.RichHandler
        show_time: true
        show_level: true
        enable_link_path: false
        level: {level.upper()}
    root:
      handlers: [{name}]
    disable_existing_loggers: false
    ...
    """

get_file_logger(name=None, level='INFO', rank_zero_only=True, fname=None)

Create a file-backed logger, optionally emitting only on rank zero.

Source code in src/ezpz/log/__init__.py

def get_file_logger(
    name: Optional[str] = None,
    level: str = "INFO",
    rank_zero_only: bool = True,
    fname: Optional[str] = None,
    # rich_stdout: bool = True,
) -> logging.Logger:
    """Create a file-backed logger, optionally emitting only on rank zero."""
    # logging.basicConfig(stream=DummyTqdmFile(sys.stderr))
    import logging

    from ezpz.dist import get_rank

    fname = "output" if fname is None else fname
    log = logging.getLogger(name)
    if rank_zero_only:
        fh = logging.FileHandler(f"{fname}.log")
        if get_rank() == 0:
            log.setLevel(level)
            fh.setLevel(level)
        else:
            log.setLevel("CRITICAL")
            fh.setLevel("CRITICAL")
    else:
        fh = logging.FileHandler(f"{fname}-{get_rank()}.log")
        log.setLevel(level)
        fh.setLevel(level)
    # create formatter and add it to the handlers
    formatter = logging.Formatter(
        "[%(asctime)s][%(name)s][%(levelname)s] - %(message)s"
    )
    fh.setFormatter(formatter)
    log.addHandler(fh)
    return log

get_logger(name=None, level=None, rank_zero_only=True, rank=None, colored_logs=True)

Return a logger initialised with the project's logging configuration.

Source code in src/ezpz/log/__init__.py

def get_logger(
    name: Optional[str] = None,
    level: Optional[str] = None,
    rank_zero_only: bool = True,
    rank: Optional[int | str] = None,
    colored_logs: Optional[bool] = True,
) -> logging.Logger:
    """Return a logger initialised with the project's logging configuration."""
    if rank is None and rank_zero_only:
        from ezpz.dist import get_rank

        rank = get_rank()
    assert rank is not None
    # if is_interactive():
    #     return get_rich_logger(name=name, level=level)
    ezpz_log_level = (
        os.environ.get("EZPZ_LOG_LEVEL", os.environ.get("LOG_LEVEL", "INFO"))
        if level is None
        else level
    )
    # level = os.environ.get("LOG_LEVEL", "INFO") if level is None else level
    # if colored_logs and use_colored_logs():
    if not colored_logs:
        os.environ["NO_COLOR"] = "1"
    logging.config.dictConfig(get_logging_config())
    logger = logging.getLogger(name if name is not None else __name__)
    if rank_zero_only:
        if int(rank) == 0:
            logger.setLevel(ezpz_log_level)
        else:
            logger.setLevel("CRITICAL")
    else:
        logger.setLevel(ezpz_log_level)
    return logger

get_logger1(name=None, level='INFO', rank_zero_only=True, **kwargs)

Legacy helper retained for compatibility; prefer :func:get_logger.

Source code in src/ezpz/log/__init__.py

def get_logger1(
    name: Optional[str] = None,
    level: str = "INFO",
    rank_zero_only: bool = True,
    **kwargs,
) -> logging.Logger:
    """Legacy helper retained for compatibility; prefer :func:`get_logger`."""
    from ezpz.dist import get_rank, get_world_size

    log = logging.getLogger(name)
    # from ezpz.log.handler import RichHandler
    from rich.logging import RichHandler as OriginalRichHandler

    from ezpz.log.console import get_console, is_interactive
    from ezpz.log.handler import RichHandler as EnrichHandler

    _ = (
        log.setLevel("CRITICAL")
        if (get_rank() == 0 and rank_zero_only)
        else log.setLevel(level)
    )
    # if rank_zero_only:
    #     if RANK != 0:
    #         log.setLevel('CRITICAL')
    #     else:
    #         log.setLevel(level)
    if get_rank() == 0:
        console = get_console(
            markup=True,  # (WORLD_SIZE == 1),
            redirect=(get_world_size() > 1),
            **kwargs,
        )
        # if console.is_jupyter:
        #     console.is_jupyter = False
        # log.propagate = True
        # log.handlers = []
        use_markup = get_world_size() == 1 and not is_interactive()
        log.addHandler(
            OriginalRichHandler(
                omit_repeated_times=False,
                level=level,
                console=console,
                show_time=True,
                show_level=True,
                show_path=True,
                markup=use_markup,
                enable_link_path=use_markup,
            )
        )
        log.setLevel(level)
    # if (
    #         len(log.handlers) > 1
    #         and all([i == log.handlers[0] for i in log.handlers])
    # ):
    #     log.handlers = [log.handlers[0]]
    if len(log.handlers) > 1 and all([i == log.handlers[0] for i in log.handlers]):
        log.handlers = [log.handlers[0]]
    enrich_handlers = get_active_enrich_handlers(log)
    found_handlers = 0
    if len(enrich_handlers) > 1:
        for h in log.handlers:
            if isinstance(h, EnrichHandler):
                if found_handlers > 1:
                    log.warning(
                        "More than one `EnrichHandler` in current logger: "
                        f"{log.handlers}"
                    )
                    log.removeHandler(h)
                found_handlers += 1
    if len(get_active_enrich_handlers(log)) > 1:
        log.warning(f"More than one `EnrichHandler` in current logger: {log.handlers}")
    return log

get_logger_new(name, level='INFO')

Return a logger configured solely via the Enrich YAML template.

Source code in src/ezpz/log/__init__.py

def get_logger_new(
    name: str,
    level: str = "INFO",
):
    """Return a logger configured solely via the Enrich YAML template."""
    import yaml

    config = yaml.safe_load(
        get_enrich_logging_config_as_yaml(name=name, level=level),
    )
    logging.config.dictConfig(config)
    log = logging.getLogger(name=name)
    log.setLevel(level)
    return log

get_rich_logger(name=None, level=None)

Return a logger backed by a single :class:RichHandler.

Source code in src/ezpz/log/__init__.py

def get_rich_logger(
    name: Optional[str] = None, level: Optional[str] = None
) -> logging.Logger:
    """Return a logger backed by a single :class:`RichHandler`."""
    from ezpz.dist import get_world_size
    from ezpz.log.handler import RichHandler

    level = "INFO" if level is None else level
    # log: logging.Logger = get_logger(name=name, level=level)
    log = logging.getLogger(name)
    log.handlers = []
    console = get_console(
        markup=True,
        redirect=(get_world_size() > 1),
    )
    handler = RichHandler(
        level,
        rich_tracebacks=False,
        console=console,
        show_path=False,
        enable_link_path=False,
    )
    log.handlers = [handler]
    log.setLevel(level)
    return log

make_layout(ratio=4, visible=True)

Define the layout.

Source code in src/ezpz/log/style.py

def make_layout(ratio: int = 4, visible: bool = True) -> Layout:
    """Define the layout."""
    layout = Layout(name="root", visible=visible)
    layout.split_row(
        Layout(name="main", ratio=ratio, visible=visible),
        Layout(name="footer", visible=visible),
    )
    return layout

print_config(config, resolve=True)

Prints content of DictConfig using Rich library and its tree structure.

Parameters:

Name	Type	Description	Default
config	DictConfig	Configuration composed by Hydra.	required
print_order	Sequence[str]	Determines in what order config components are printed.	required
resolve	bool	Whether to resolve reference fields of DictConfig.	True

Source code in src/ezpz/log/style.py

def print_config(
    config: DictConfig | dict | Any,
    resolve: bool = True,
) -> None:
    """Prints content of DictConfig using Rich library and its tree structure.

    Args:
        config (DictConfig): Configuration composed by Hydra.
        print_order (Sequence[str], optional): Determines in what order config
            components are printed.
        resolve (bool, optional): Whether to resolve reference fields of
            DictConfig.
    """
    import pandas as pd

    tree = rich.tree.Tree("CONFIG")  # , style=style, guide_style=style)
    quee = []
    for f in config:
        if f not in quee:
            quee.append(f)
    dconfig = {}
    for f in quee:
        branch = tree.add(f)  # , style=style, guide_style=style)
        config_group = config[f]
        if isinstance(config_group, DictConfig):
            branch_content = OmegaConf.to_yaml(config_group, resolve=resolve)
            cfg = OmegaConf.to_container(config_group, resolve=resolve)
        else:
            branch_content = str(config_group)
            cfg = str(config_group)
        dconfig[f] = cfg
        branch.add(rich.syntax.Syntax(branch_content, "yaml"))
    outfile = Path(os.getcwd()).joinpath("config_tree.log")
    from rich.console import Console

    with outfile.open("wt") as f:
        console = Console(file=f)
        console.print(tree)
    with open("config.json", "w") as f:
        f.write(json.dumps(dconfig))
    cfgfile = Path("config.yaml")
    OmegaConf.save(config, cfgfile, resolve=True)
    cfgdict = OmegaConf.to_object(config)
    logdir = Path(os.getcwd()).resolve().as_posix()
    if not config.get("debug_mode", False):
        dbfpath = Path(os.getcwd()).joinpath("logdirs.csv")
    else:
        dbfpath = Path(os.getcwd()).joinpath("logdirs-debug.csv")
    if dbfpath.is_file():
        mode = "a"
        header = False
    else:
        mode = "w"
        header = True
    df = pd.DataFrame({logdir: cfgdict})
    df.T.to_csv(dbfpath.resolve().as_posix(), mode=mode, header=header)
    os.environ["LOGDIR"] = logdir

print_styles()

Print the configured logging styles (optionally exporting to HTML).

Source code in src/ezpz/log/__init__.py

def print_styles():
    """Print the configured logging styles (optionally exporting to HTML)."""
    import argparse

    parser = argparse.ArgumentParser()
    from rich.text import Text

    from ezpz.log.console import Console

    parser.add_argument("--html", action="store_true", help="Export as HTML table")
    args = parser.parse_args()
    html: bool = args.html
    from rich.table import Table

    console = Console(record=True, width=120) if html else Console()
    table = Table("Name", "Styling")
    for style_name, style in STYLES.items():
        table.add_row(Text(style_name, style=style), str(style))

    console.print(table)
    if html:
        outfile = "enrich_styles.html"
        print(f"Saving to `{outfile}`")
        with open(outfile, "w") as f:
            f.write(console.export_html(inline_styles=True))

print_styles_alt(html=False, txt=False)

Variant of :func:print_styles with HTML and plain-text exports.

Source code in src/ezpz/log/__init__.py

def print_styles_alt(
    html: bool = False,
    txt: bool = False,
):
    """Variant of :func:`print_styles` with HTML and plain-text exports."""
    from pathlib import Path

    from rich.table import Table
    from rich.text import Text

    from ezpz.log.console import get_console
    from ezpz.log.style import DEFAULT_STYLES

    console = get_console(record=html, width=150)
    table = Table("Name", "Styling")
    styles = DEFAULT_STYLES
    styles |= STYLES
    for style_name, style in styles.items():
        table.add_row(Text(style_name, style=style), str(style))
    console.print(table)
    if html:
        outfile = "ezpz_styles.html"
        print(f"Saving to `{outfile}`")
        with open(outfile, "w") as f:
            f.write(console.export_html(inline_styles=True))
    if txt:
        file1 = "ezpz_styles.txt"
        text = console.export_text()
        # with open(file1, "w") as file:
        with Path(file1).open("w") as file:
            file.write(text)

printarr(*arrs, float_width=6)

Print a pretty table giving name, shape, dtype, type, and content information for input tensors or scalars.

Call like: printarr(my_arr, some_other_arr, maybe_a_scalar). Accepts a variable number of arguments.

Inputs can be

• Numpy tensor arrays
• Pytorch tensor arrays
• Jax tensor arrays
• Python ints / floats
• None

It may also work with other array-like types, but they have not been tested

Use the float_width option specify the precision to which floating point types are printed.

Author: Nicholas Sharp (nmwsharp.com) Canonical source: https://gist.github.com/nmwsharp/54d04af87872a4988809f128e1a1d233 License: This snippet may be used under an MIT license, and it is also released into the public domain. Please retain this docstring as a reference.

Source code in src/ezpz/log/style.py

def printarr(*arrs, float_width=6):
    """
    Print a pretty table giving name, shape, dtype, type, and content
    information for input tensors or scalars.

    Call like: printarr(my_arr, some_other_arr, maybe_a_scalar). Accepts a
    variable number of arguments.

    Inputs can be:
        - Numpy tensor arrays
        - Pytorch tensor arrays
        - Jax tensor arrays
        - Python ints / floats
        - None

    It may also work with other array-like types, but they have not been tested

    Use the `float_width` option specify the precision to which floating point
    types are printed.

    Author: Nicholas Sharp (nmwsharp.com)
    Canonical source:
        https://gist.github.com/nmwsharp/54d04af87872a4988809f128e1a1d233
    License: This snippet may be used under an MIT license, and it is also
    released into the public domain. Please retain this docstring as a
    reference.
    """
    import inspect

    frame_ = inspect.currentframe()
    assert frame_ is not None
    frame = frame_.f_back
    # if frame_ is not None:
    #     frame = frame_.f_back
    # else:
    #     frame = inspect.getouterframes()
    default_name = "[temporary]"

    # helpers to gather data about each array

    def name_from_outer_scope(a):
        if a is None:
            return "[None]"
        name = default_name
        if frame_ is not None:
            for k, v in frame_.f_locals.items():
                if v is a:
                    name = k
                    break
        return name

    def dtype_str(a):
        if a is None:
            return "None"
        if isinstance(a, int):
            return "int"
        if isinstance(a, float):
            return "float"
        return str(a.dtype)

    def shape_str(a):
        if a is None:
            return "N/A"
        if isinstance(a, int):
            return "scalar"
        if isinstance(a, float):
            return "scalar"
        return str(list(a.shape))

    def type_str(a):
        # TODO this is is weird... what's the better way?
        return str(type(a))[8:-2]

    def device_str(a):
        if hasattr(a, "device"):
            device_str = str(a.device)
            if len(device_str) < 10:
                # heuristic: jax returns some goofy long string we don't want,
                # ignore it
                return device_str
        return ""

    def format_float(x):
        return f"{x:{float_width}g}"

    def minmaxmean_str(a):
        if a is None:
            return ("N/A", "N/A", "N/A")
        if isinstance(a, int) or isinstance(a, float):
            return (format_float(a), format_float(a), format_float(a))

        # compute min/max/mean. if anything goes wrong, just print 'N/A'
        min_str = "N/A"
        try:
            min_str = format_float(a.min())
        except Exception:
            pass
        max_str = "N/A"
        try:
            max_str = format_float(a.max())
        except Exception:
            pass
        mean_str = "N/A"
        try:
            mean_str = format_float(a.mean())
        except Exception:
            pass

        return (min_str, max_str, mean_str)

    try:
        props = [
            "name",
            "dtype",
            "shape",
            "type",
            "device",
            "min",
            "max",
            "mean",
        ]

        # precompute all of the properties for each input
        str_props = []
        for a in arrs:
            minmaxmean = minmaxmean_str(a)
            str_props.append(
                {
                    "name": name_from_outer_scope(a),
                    "dtype": dtype_str(a),
                    "shape": shape_str(a),
                    "type": type_str(a),
                    "device": device_str(a),
                    "min": minmaxmean[0],
                    "max": minmaxmean[1],
                    "mean": minmaxmean[2],
                }
            )

        # for each property, compute its length
        maxlen = {}
        for p in props:
            maxlen[p] = 0
        for sp in str_props:
            for p in props:
                maxlen[p] = max(maxlen[p], len(sp[p]))

        # if any property got all empty strings,
        # don't bother printing it, remove if from the list
        props = [p for p in props if maxlen[p] > 0]

        # print a header
        header_str = ""
        for p in props:
            prefix = "" if p == "name" else " | "
            fmt_key = ">" if p == "name" else "<"
            header_str += f"{prefix}{p:{fmt_key}{maxlen[p]}}"
        print(header_str)
        print("-" * len(header_str))
        # now print the acual arrays
        for strp in str_props:
            for p in props:
                prefix = "" if p == "name" else " | "
                fmt_key = ">" if p == "name" else "<"
                print(f"{prefix}{strp[p]:{fmt_key}{maxlen[p]}}", end="")
            print("")

    finally:
        del frame

should_do_markup(stream=sys.stdout)

Decide about use of ANSI colors.

Source code in src/ezpz/log/console.py

def should_do_markup(stream: TextIO = sys.stdout) -> bool:
    """Decide about use of ANSI colors."""
    py_colors = None

    # https://xkcd.com/927/
    for env_var in [
        "PY_COLORS",
        "CLICOLOR",
        "FORCE_COLOR",
        "ANSIBLE_FORCE_COLOR",
    ]:
        value = os.environ.get(env_var, None)
        if value is not None:
            py_colors = to_bool(value)
            break

    # If deliverately disabled colors
    if os.environ.get("NO_COLOR", None):
        return False

    # User configuration requested colors
    if py_colors is not None:
        return to_bool(py_colors)

    term = os.environ.get("TERM", "")
    if "xterm" in term:
        return True

    if term.lower() == "dumb":
        return False

    # Use tty detection logic as last resort.
    # Because there are numerous factors that can make isatty return a
    # misleading value, including:
    # - stdin.isatty() is the only one returning true, even on a real terminal
    # - stderr returning false if user uses an error stream coloring solution
    return stream.isatty()

to_bool(value)

Return a bool for the arg.

Source code in src/ezpz/log/console.py

def to_bool(value: Any) -> bool:
    """Return a bool for the arg."""
    if value is None or isinstance(value, bool):
        return bool(value)
    if isinstance(value, str):
        value = value.lower()
    if value in ("yes", "on", "1", "true", 1):
        return True
    return False

Usage Examples

Create a Rank-Aware Logger

import ezpz.log as ezlog

logger = ezlog.get_logger("train")
logger.info("hello from rank 0")

Print Available Styles

from ezpz.log import print_styles

print_styles()