Altering Python Function Signature Typehints

Published
6 min read

I've recently worked on a piece of asyncio Python software, which needed to leverage a synchronous API client library. Implementing synchronous code in asyncio Python is easy enough, and can be done using an async Executor to run the synchronous code in a separate thread or process, avoiding it blocking the main thread. While implementing this part of the project, I found myself not liking that with the raw proxy, we would be await-ing methods that were typehinted as not being coroutines, so decided to figure out how we could achieve this, without needing to maintain multiple versions of the API client, or manually write async versions of the methods we needed.

All the code for this post is also available here on my Github account.

I'll start by explaining a simplified version how the async API client's codebase was laid out, and then go on to explain how the automated typehint alteration was implemented on top of this. Let's assume a very basic synchronous API client as so:

sync_client_lib/client.py

import requests

from typing import Any

class ApiEndpointCollection:
    def retrieve_resource(self, resource_id: int) -> dict[str, Any]:
        """
        Retrieve a resource from the API

        :param resource_id: The ID of the resource to retrieve
        """
        return requests.get(f"https://httpbin.org/get?id={resource_id}").json()


class ApiClient:
    endpoints: ApiEndpointCollection

    def __init__(self):
        self.endpoints = ApiEndpointCollection()

And the async wrapper proxy looks like this:

api_client_async.py

import asyncio
import functools

from sync_client_lib.client import ApiClient, ApiEndpointCollection

class AsyncApiEndpointProxy:
    _method: callable

    def __init__(self, method: callable):
        self._method = method

    async def __call__(self, *args, **kwargs):
        return await asyncio.get_event_loop().run_in_executor(
            None, functools.partial(self._method, *args, **kwargs)
        )


class AsyncApiEndpointCollectionProxy:
    _collection: ApiEndpointCollection

    def __init__(self, collection: ApiEndpointCollection):
        self._collection = collection

    def __getattr__(self, name: str):
        return AsyncApiEndpointProxy(getattr(self._collection, name))


class AsyncApiClientProxy:
    _client: ApiClient

    endpoints: ApiEndpointCollection

    def __init__(self):
        self._client = ApiClient()

    def __getattr__(self, name: str):
        return AsyncApiEndpointCollectionProxy(getattr(self._client, name))


if __name__ == "__main__":
    client = AsyncApiClientProxy()

    print(asyncio.run(client.endpoints.retrieve_resource(42)))

Here's the typehint that this produces:

Python Async Proxy - Bad Typehint

This achieves the goal of transparently proxying methods to the synchronous library, and allowing a typehint for all the methods to be seen by the caller, but also has the less than ideal side-effect of the typehinting all being for the synchronous library. To resolve this, I had to figure out an alternative approach that provided maximum flexibility and the minimum code changes required.

After some research, I decided to use a combination of decorators and some advanced usage of Python's built in typehinting library, and also separating the sync and async methods, to allow both to be typehinted correctly. The code became the following, which I'll explain below:

api_client_async.py

import asyncio
import functools

from typing import Any, Callable, Coroutine, ParamSpec, TypeVar

from sync_client_lib.client import ApiClient, ApiEndpointCollection

P = ParamSpec("P")
R = TypeVar("R")

SyncFunc = Callable[P, R]
CoroFunc = Callable[P, Coroutine[Any, Any, R]]
CoroRetFunc = Callable[P, Coroutine[Any, Any, R]]
CoroFuncWrapper = Callable[[SyncFunc[P, R]], CoroFunc[P, R]]


def _typehint_coroutine() -> CoroFuncWrapper[P, R]:
    """Decorate a function's typehint to return its normal value wrapped in a coro."""

    def outer(fn: SyncFunc[P, R]) -> CoroRetFunc[P, R]:
        @functools.wraps(fn)
        def wrapper(*args: P.args, **kwargs: P.kwargs):
            return fn(*args, **kwargs)

        return wrapper

    return outer


class AsyncApiEndpointProxy:
    _method: Callable
    _convert_async: bool

    def __init__(self, method: Callable, convert_async: bool):
        self._method = method
        self._convert_async = convert_async

    async def __call__(self, *args, **kwargs):
        if self._convert_async:
            return await asyncio.get_event_loop().run_in_executor(
                None, functools.partial(self._method, *args, **kwargs)
            )
        else:
            return self._method(*args, **kwargs)


class AsyncApiEndpointCollectionProxy:
    _collection: ApiEndpointCollection

    def __init__(self, collection: ApiEndpointCollection):
        self._collection = collection

    def __getattr__(self, name: str):
        convert_async = name.endswith("_async")

        return AsyncApiEndpointProxy(
            getattr(
                self._collection,
                name[:-6] if convert_async else name,
                convert_async,
            )
        )


class AsyncApiEndpointCollection(ApiEndpointCollection):
    retrieve_resource_async = _typehint_coroutine()(
        ApiEndpointCollection.retrieve_resource
    )


class AsyncApiClientProxy:
    _client: ApiClient

    endpoints: AsyncApiEndpointCollection

    def __init__(self):
        self._client = ApiClient()

    def __getattr__(self, name: str):
        return AsyncApiEndpointCollectionProxy(getattr(self._client, name))


if __name__ == "__main__":
    client = AsyncApiClientProxy()

    print(asyncio.run(client.endpoints.retrieve_resource_async(43)))

The basic approach is to subclass the endpoint collection, and add an automatically converted typehint for async variants of any methods we're going to use. The main downside here is that we'll need to add any methods that we want to use in an asyncio context to the endpoint collection subclass, but the benefit is that if any of these methods change in the synchronous API client implementation, they will continue to be typehinted correctly. As this application (and most that I've worked on) use a relatively limited set of API endpoint calls that aren't constantly changing, and this solution does still provide full typehinting for all the synchronous API client methods, this was an acceptable tradeoff.

The function that does the typehint modification uses some advanced typehinting techniques, so I'll explain how it works from start to finish. Here's all the important code for it in one place:

P = ParamSpec("P")
R = TypeVar("R")

SyncFunc = Callable[P, R]
CoroFunc = Callable[P, Coroutine[Any, Any, R]]
CoroRetFunc = Callable[P, Coroutine[Any, Any, R]]
CoroFuncWrapper = Callable[[SyncFunc[P, R]], CoroFunc[P, R]]


def _typehint_coroutine() -> CoroFuncWrapper[P, R]:
    """Decorate a function's typehint to return its normal value wrapped in a coro."""

    def outer(fn: SyncFunc[P, R]) -> CoroRetFunc[P, R]:
        @functools.wraps(fn)
        def wrapper(*args: P.args, **kwargs: P.kwargs):
            return fn(*args, **kwargs)

        return wrapper

    return outer

At its core, it's just a simple function wrapper decorator, the same as are regularly used in many Python applications and libraries. The fancy footwork is in the use of typing.ParamSpec and typing.TypeVar. TypeVar is commonly used when writing generic classes, and allows us to typehint individual dynamic types. ParamSpec, defined in PEP-0612 is a more recent addition to Python, and only became part of the core typing library since Python 3.10. It allows us to define generic function parameters collections, the same way as we would typehint single parameters and return values with TypeVar.

We've defined various type aliases above the function definition, in order to make the function a bit more readable. It may appear at first glance that there is unnecessary redundancy, and duplicate definitions, but every part of the code is critical to correct operation.

SyncFunc describes the synchronous function that we're wrapping. You can see that we've defined it as a Callable, with dynamic arguments and return value.
CoroFunc is the typehint for the synchronous function, if it were a coroutine.
CoroRetFunc has the same signature as CoroFunc, and does refer to the same function, but is used in a different location in the decorator chain, so must have a separate definition in order for the typehints to be correctly calculated by the IDE.
CoroFuncWrapper describes the return value of the decorator itself.

The result of this is that now, we have typehints for both the sync ...

Python Async Proxy - Good Synchronous Typehint

... and async variant of the method:

Python Async Proxy - Good Asynchronous Typehint

While almost certainly somewhat over-engineered for the relatively minimal effect, this was nevertheless a super fun challenge. Until next time!

Alex