File size: 4,419 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
import inspect
from functools import partial
from typing import (
Any,
Callable,
Iterable,
List,
Optional,
Tuple,
Type,
TypeVar,
Union,
overload,
)
T = TypeVar("T")
Result = Iterable[Union[Any, Tuple[Any], Tuple[str, Any], Tuple[str, Any, Any]]]
RichReprResult = Result
class ReprError(Exception):
"""An error occurred when attempting to build a repr."""
@overload
def auto(cls: Optional[Type[T]]) -> Type[T]:
...
@overload
def auto(*, angular: bool = False) -> Callable[[Type[T]], Type[T]]:
...
def auto(
cls: Optional[Type[T]] = None, *, angular: Optional[bool] = None
) -> Union[Type[T], Callable[[Type[T]], Type[T]]]:
"""Class decorator to create __repr__ from __rich_repr__"""
def do_replace(cls: Type[T], angular: Optional[bool] = None) -> Type[T]:
def auto_repr(self: T) -> str:
"""Create repr string from __rich_repr__"""
repr_str: List[str] = []
append = repr_str.append
angular: bool = getattr(self.__rich_repr__, "angular", False) # type: ignore[attr-defined]
for arg in self.__rich_repr__(): # type: ignore[attr-defined]
if isinstance(arg, tuple):
if len(arg) == 1:
append(repr(arg[0]))
else:
key, value, *default = arg
if key is None:
append(repr(value))
else:
if default and default[0] == value:
continue
append(f"{key}={value!r}")
else:
append(repr(arg))
if angular:
return f"<{self.__class__.__name__} {' '.join(repr_str)}>"
else:
return f"{self.__class__.__name__}({', '.join(repr_str)})"
def auto_rich_repr(self: Type[T]) -> Result:
"""Auto generate __rich_rep__ from signature of __init__"""
try:
signature = inspect.signature(self.__init__)
for name, param in signature.parameters.items():
if param.kind == param.POSITIONAL_ONLY:
yield getattr(self, name)
elif param.kind in (
param.POSITIONAL_OR_KEYWORD,
param.KEYWORD_ONLY,
):
if param.default is param.empty:
yield getattr(self, param.name)
else:
yield param.name, getattr(self, param.name), param.default
except Exception as error:
raise ReprError(
f"Failed to auto generate __rich_repr__; {error}"
) from None
if not hasattr(cls, "__rich_repr__"):
auto_rich_repr.__doc__ = "Build a rich repr"
cls.__rich_repr__ = auto_rich_repr # type: ignore[attr-defined]
auto_repr.__doc__ = "Return repr(self)"
cls.__repr__ = auto_repr # type: ignore[assignment]
if angular is not None:
cls.__rich_repr__.angular = angular # type: ignore[attr-defined]
return cls
if cls is None:
return partial(do_replace, angular=angular)
else:
return do_replace(cls, angular=angular)
@overload
def rich_repr(cls: Optional[Type[T]]) -> Type[T]:
...
@overload
def rich_repr(*, angular: bool = False) -> Callable[[Type[T]], Type[T]]:
...
def rich_repr(
cls: Optional[Type[T]] = None, *, angular: bool = False
) -> Union[Type[T], Callable[[Type[T]], Type[T]]]:
if cls is None:
return auto(angular=angular)
else:
return auto(cls)
if __name__ == "__main__":
@auto
class Foo:
def __rich_repr__(self) -> Result:
yield "foo"
yield "bar", {"shopping": ["eggs", "ham", "pineapple"]}
yield "buy", "hand sanitizer"
foo = Foo()
from rich.console import Console
console = Console()
console.rule("Standard repr")
console.print(foo)
console.print(foo, width=60)
console.print(foo, width=30)
console.rule("Angular repr")
Foo.__rich_repr__.angular = True # type: ignore[attr-defined]
console.print(foo)
console.print(foo, width=60)
console.print(foo, width=30)
|