diff --git a/src/runboat/models.py b/src/runboat/models.py index 9c5d0fe..7a2c483 100644 --- a/src/runboat/models.py +++ b/src/runboat/models.py @@ -156,10 +156,10 @@ class Build(BaseModel): def live_link(self) -> str: return f"{self.webui_link}?live" - async def init_log(self) -> str: + async def init_log(self) -> str | None: return await k8s.log(self.name, job_kind=k8s.DeploymentMode.initialize) - async def log(self) -> str: + async def log(self) -> str | None: return await k8s.log(self.name, job_kind=None) @classmethod diff --git a/src/runboat/utils.py b/src/runboat/utils.py index 38ebd31..95f28cc 100644 --- a/src/runboat/utils.py +++ b/src/runboat/utils.py @@ -3,6 +3,16 @@ import functools import re from concurrent.futures.thread import ThreadPoolExecutor from functools import wraps +from typing import ( + Any, + AsyncGenerator, + Awaitable, + Callable, + Generator, + Iterator, + ParamSpec, + TypeVar, +) _pool = ThreadPoolExecutor(max_workers=20, thread_name_prefix="sync_to_async") @@ -11,29 +21,38 @@ def slugify(s: str | int) -> str: return re.sub(r"[^a-z0-9]", "-", str(s).lower()) -def sync_to_async(func): +# TODO replace ... with P below when mypy supports PEP 612 +# (https://github.com/python/mypy/issues/8645) +P = ParamSpec("P") +R = TypeVar("R") +T = TypeVar("T") + + +def sync_to_async(func: Callable[..., R]) -> Callable[..., Awaitable[R]]: @wraps(func) - async def inner(*args, **kwargs): + async def inner(*args: Any, **kwargs: Any) -> R: f = functools.partial(func, *args, **kwargs) return await asyncio.get_running_loop().run_in_executor(_pool, f) return inner -def sync_to_async_iterator(iterator_func): +def sync_to_async_iterator( + iterator_func: Callable[..., Generator[R, None, None]] +) -> Callable[..., AsyncGenerator[R, None]]: @sync_to_async - def async_next(iterator): + def async_next(iterator: Iterator[R]) -> R: try: return next(iterator) except StopIteration: raise StopAsyncIteration() @sync_to_async - def async_iterator_func(*args, **kwargs): + def async_iterator_func(*args: Any, **kwargs: Any) -> Generator[R, None, None]: return iterator_func(*args, **kwargs) @wraps(iterator_func) - async def inner(*args, **kwargs): + async def inner(*args: Any, **kwargs: Any) -> AsyncGenerator[R, None]: iterator = await async_iterator_func(*args, **kwargs) while True: try: