123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149 |
- 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 == 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 pip._vendor.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)
|