Python annotations for decorated async functions

Question

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]?

Answer 1

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