__all__ = [
"repr_short",
"repr_type",
"repr_source_path",
"repr_pyfunction",
"rich_repr",
]
import inspect
import os
import sys
import types
import typing
from collections.abc import Callable, Iterable
from functools import partial
from typing import Any, TypeVar, overload
from typing_extensions import TypeAliasType
import rich
from rich.color import Color
from rich.style import Style
from rich.text import Text
from ._alias import _transform_union_alias
T = TypeVar("T")
path_style = Style(color=Color.from_ansi(7))
file_style = Style(bold=True, underline=True)
module_style = Style(color="grey66")
class_style = Style(bold=True)
[docs]
def repr_type(x: object, /) -> Text:
"""Returns a :class:`rich.Text` representation of a type or module.
Does some light syntax highlighting mimicking Julia, boldening class names and
coloring module names with a lighter color.
Args:
x (object): Type or module.
Returns:
:class:`rich.Text`: Representation.
"""
# Apply union aliasing if `x` is a union. This allows us to have the correct
# syntax highlighting for aliased unions.
x = _transform_union_alias(x)
if isinstance(x, type):
if x.__module__ in ["builtins", "typing", "typing_extensions"]:
return Text(x.__qualname__, style=class_style)
else:
res = Text(f"{x.__module__}.", style=module_style)
res += Text(x.__qualname__, style=class_style)
return res
if isinstance(x, types.ModuleType):
return Text(x.__name__, style=module_style)
return Text(repr(x), style=class_style)
[docs]
def repr_short(x: object, /) -> str:
"""Representation as a string, but in shorter form. This just calls
:func:`typing._type_repr`.
If the type is a union registered in Plum's alias registry, the alias name
is used instead.
Args:
x (object): Object.
Returns:
str: Shorter representation of `x`.
"""
if isinstance(transformed := _transform_union_alias(x), TypeAliasType):
# It's an aliased union. Use the alias name.
return str(transformed.__name__)
# :func:`typing._type_repr` is an internal function, but it should be
# available in Python versions 3.9 through 3.14.
return typing._type_repr(x)
def _safe_getfile(obj: object, /) -> str:
"""Safer version of :func:`inspect.getfile`.
Classes defined in PyBind, even if they pretend hard to be functions, like
`jax.jit(fun)`, will raise an error if passed to `inspect.getfile`.
This function will catch those errors and try harder to get the underlying file, and
raise an error only if it cannot.
This function can only raise an `OSError`.
Args:
obj (object): An object.
Returns:
str: Path to the file that defines `obj`.
Raises:
OSError: File that defines `obj` cannot be found.
"""
try:
return inspect.getfile(obj) # type: ignore[arg-type]
except TypeError:
# Raised when the function passed is a C-defined class. It might still contain
# `__module__`, which can be used to backtrace the file that defines `obj`.
if hasattr(obj, "__module__"):
module = sys.modules.get(obj.__module__, None)
if module is not None and hasattr(module, "__file__"):
file_path = module.__file__
if file_path is not None:
return file_path
raise OSError("Source code not available.") from None
[docs]
def repr_source_path(function: Callable[..., Any], /) -> Text | None:
"""Create a :class:`rich.Text` object with an hyperlink to the function definition.
Args:
function (Callable[..., Any]): A function.
Returns:
:class:`rich.Text` or None:
Representation with a hyperlink to the source. If the introspection
failed, it returns :obj:`None`.
"""
try:
f_path = _safe_getfile(function)
f_line = inspect.getsourcelines(function)[1]
uri = f"file://{f_path}#{f_line}"
# Compress the path.
home_path = os.path.expanduser("~")
f_path = f_path.replace(home_path, "~")
# Underline file name.
f_name = os.path.basename(f_path)
f_path = (
Text(os.path.dirname(f_path), style=path_style)
+ Text("/")
+ Text(f_name, style=file_style)
)
f_path.append_text(Text(f":{f_line}"))
f_path.stylize(f"link {uri}")
return f_path
except OSError: # pragma: no cover
return None
[docs]
def repr_pyfunction(f: Callable[..., Any]) -> Text:
"""Create a :class:`rich.Text` object representation a function, including a link
to the source definition created with :func:`repr_source_path`.
Args:
f (Callable[..., Any]): A function.
Returns:
:class:`rich.Text`: Representation of `f`.
"""
res = Text(repr(f))
source = repr_source_path(f)
if source is not None:
res.append(" @ ")
res.append_text(source)
return res
########################
# Rich class decorator #
########################
def __repr_from_rich__(self: object, /) -> str:
"""Default implementation of `__repr__` that calls :mod:`rich`.
Returns:
str: Representation of `self.`
"""
console = rich.get_console()
with console.capture() as capture:
console.print(self, end="")
res: str = capture.get()
return res
def _repr_mimebundle_from_rich_(
self: object, include: Iterable[str], exclude: Iterable[str], **kw_args: Any
) -> dict[str, str]:
"""Implementation of `_repr_mimebundle_` for better rendering in Jupyter.
Args:
include (Iterable[str]): Only these MIME types should be included.
exclude (Iterable[str]): These MIME types should be excluded.
**kw_args (object): Additional keyword arguments. These are ignored.
Returns:
dict[str, str]: Representation by MIME type.
"""
from rich.jupyter import _render_segments
console = rich.get_console()
segments = list(console.render(self, console.options))
html = _render_segments(segments)
text = console._render_buffer(segments)
data = {"text/plain": text, "text/html": html}
if include:
data = {k: v for (k, v) in data.items() if k in include}
if exclude:
data = {k: v for (k, v) in data.items() if k not in exclude}
return data
@overload
def rich_repr(cls: type[T], str: bool = False) -> type[T]: ...
@overload
def rich_repr(cls: None = None, str: bool = False) -> Callable[[type[T]], type[T]]: ...
[docs]
def rich_repr(
cls: type[T] | None = None, str: bool = False
) -> type[T] | Callable[[type[T]], type[T]]:
"""Class decorator defining a `__repr__` method that calls :mod:`rich.`
This also sets `_repr_mimebundle_` for better rendering in Jupyter.
Args:
cls (type, optional): Class to decorate. If `None`, this function returns a
decorator.
str (bool, optional): Also define `__str__`. Defaults to not defining
`__str__`
Returns:
The decorated class. If cls is None, returns a decorator.
"""
if cls is None:
return partial(rich_repr, str=str)
cls.__repr__ = __repr_from_rich__ # type: ignore[method-assign]
cls._repr_mimebundle_ = _repr_mimebundle_from_rich_ # type: ignore[attr-defined]
if str:
cls.__str__ = __repr_from_rich__ # type: ignore[method-assign]
return cls
