repr.py 4.3 KB

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