From 4c2bc266b3bd3a01faf7a6eca855911c91b2701e Mon Sep 17 00:00:00 2001 From: James Braza Date: Fri, 6 Dec 2024 10:17:42 -0800 Subject: [PATCH] Type hinted `maybe_wait_for` return (#143) --- src/aviary/env.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/aviary/env.py b/src/aviary/env.py index b8be8752..9e9d29af 100644 --- a/src/aviary/env.py +++ b/src/aviary/env.py @@ -9,7 +9,7 @@ from abc import ABC, abstractmethod from collections.abc import Awaitable, Iterator from copy import deepcopy -from typing import Annotated, Any, Generic, Self, TypeAlias, TypeVar, cast +from typing import Annotated, Generic, Self, TypeAlias, TypeVar, cast from pydantic import ( BaseModel, @@ -39,7 +39,10 @@ Serializable: TypeAlias = dict | list | int | float | str | bool | BaseModel -async def maybe_wait_for(future: Awaitable, timeout: float | None) -> Any: +_T = TypeVar("_T") + + +async def maybe_wait_for(future: Awaitable[_T], timeout: float | None) -> _T: """Apply a timeout to an awaitable if one is provided, otherwise await indefinitely.""" if timeout is None: return await future