2024-10-04 05:08:30 +00:00
|
|
|
from typing import Callable, Optional, Union
|
2024-07-08 17:15:05 +00:00
|
|
|
|
2024-07-07 14:13:55 +00:00
|
|
|
from rich.console import ConsoleRenderable, Group, RichCast
|
|
|
|
from rich.progress import Progress
|
|
|
|
|
|
|
|
|
|
|
|
class CustomProgress(Progress):
|
2024-09-01 06:27:53 +00:00
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
*args,
|
2024-10-04 05:08:30 +00:00
|
|
|
cust_objs: Optional[list[ConsoleRenderable]] = None,
|
|
|
|
cust_callables: Optional[list[Callable[[], ConsoleRenderable]]] = None,
|
2024-09-01 06:27:53 +00:00
|
|
|
**kwargs,
|
|
|
|
) -> None:
|
|
|
|
self._cust_objs = cust_objs or []
|
|
|
|
self._cust_callables = cust_callables or []
|
|
|
|
|
2024-07-07 14:13:55 +00:00
|
|
|
super().__init__(*args, **kwargs)
|
|
|
|
|
2024-07-08 17:15:05 +00:00
|
|
|
def get_renderable(self) -> Union[ConsoleRenderable, RichCast, str]:
|
2024-08-29 05:08:41 +00:00
|
|
|
objs = [obj for obj in self._cust_objs]
|
|
|
|
for cust_call in self._cust_callables:
|
|
|
|
objs.append(cust_call())
|
|
|
|
renderable = Group(*objs, *self.get_renderables())
|
2024-07-07 14:13:55 +00:00
|
|
|
return renderable
|