I have difficulties with annotations for my coroutines which are decorated to prevent aiohttp errors. There are my two functions:
from typing import Callable, Awaitable, Optional
from os import sep
import aiofiles
import aiohttp
from asyncio.exceptions import TimeoutError
from aiohttp.client_exceptions import ClientError
def catch_aiohttp_errors(func: Callable[..., Awaitable]) -> Callable[..., Awaitable]:
async def wrapper(*args):
try:
return await func(*args)
except (TimeoutError, ClientError):
return None
return wrapper
@catch_aiohttp_errors
async def download(url: str, download_path: str, filename: str, suffix: str) -> Optional[str]:
async with aiohttp.ClientSession() as session:
async with session.get(url) as response:
async with aiofiles.open(download_path sep filename '.' suffix, 'wb') as file:
async for chunk in response.content.iter_chunked(1024):
await file.write(chunk) if chunk else await file.write(b'')
return download_path sep filename '.' suffix
The main reason to make decorator function is that i have several async functions using aiohttp, and i don't want to write try/except
statements in every similar function.
The problem i faced is correct annotation for my second function.
As you can see, it returns str
. But if there will be errors, it will return None
according to try/except
part of the decorator function.
Is it correct to annotate such function with Optional[str]
?
CodePudding user response:
I would suggest using TypeVar
as Awaitable
type parameter to stop losing information about decorated function: in your example result of call to download
would be of type Any
. Also using ParamSpec
will help preserve arguments. Finally, something like this should work (assuming python 3.10, replace all unknown typing
imports with typing_extensions
otherwise):
from typing import Callable, Awaitable, Optional, TypeVar, ParamSpec
from functools import wraps
_T = TypeVar('_T')
_P = ParamSpec('_P')
def catch_aiohttp_errors(func: Callable[_P, Awaitable[_T]]) -> Callable[_P, Awaitable[Optional[_T]]]:
@wraps(func)
async def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> Optional[_T]:
try:
return await func(*args)
except Exception:
return None
return wrapper
@catch_aiohttp_errors
async def download(url: str, download_path: str, filename: str, suffix: str) -> str:
return 'foo'
Now download
has signature
def (url: builtins.str, download_path: builtins.str, filename: builtins.str, suffix: builtins.str) -> typing.Awaitable[Union[builtins.str, None]]
Also you don't have to add Optional
manually now - decorator will do. Playground with this solution