Type hinting a Python contextmanager with arguments

user_315168 0
Viewed 34

I have the following context manager:

@contextmanager
def timed_task(task_name: str, **context_info)
     pass

Now I have the function that

def my_func(my_timed_task: Any):
     with my_timed_task("my_func_task", foo="bar"):
          pass

How would I type hint my_func so that it knows my_timed_task is a context manager of type timed_task(task_name: str, **kwargs), or any equivalent context manager with the same arguments?

I know about contextlib.AbstractContextManager but I could find documentation that explains how to use it with the conjunction of contextmanager arguments, instead of just argumentless context managers.

contextmanager type-hinting python
1 Answers

So I am making a few assumptions here to provide a complete answer. If any of them are incorrect, you will have to adjust the code accordingly. I assume:

  • Your timed_task and my_func both can return any type.
  • Your **context_info kwargs can be any type.

One thing that is important to understand here, is that your decorated function itself is not the context manager. It is a factory function that returns a context manager, i.e. an object that has its __enter__ and __exit__ methods defined for usage in a with-statement. (see docs)

In the matter at hand, you basically have two choices.

The first one that is technically the "most correct" in my opinion, but may also be overkill in your situation as well as in many other situations, is defining your own Protocol first:

from contextlib import contextmanager, AbstractContextManager
from typing import Any, Protocol


class MyContextManagerFactory(Protocol):
    def __call__(self, task_name: str, **context_info: Any) -> AbstractContextManager[Any]: ...


@contextmanager
def timed_task(task_name: str, **context_info: Any) -> Any:
    pass


def my_func(my_timed_task: MyContextManagerFactory) -> Any:
    with my_timed_task("my_func_task", foo="bar"):
        pass

This has mypy pass (in --strict mode) without a problem.

The reason a protocol is needed here, is because you have arbitrary keyword arguments **context_info allowed in your context manager factory, and to my knowledge it is currently not possible to specify a Callable type accordingly.

The benefit of being this pedantic with your types is that an IDE like PyCharm will give you precise hints not only about what arguments of what type are allowed when calling my_timed_task, but actually about the full signature including the parameter names, which is neat.

Your second option is far simpler, but also less precise. You can simply define the type of my_timed_task as a callable accepting anything and returning your context manager:

from contextlib import contextmanager, AbstractContextManager
from typing import Any, Callable


@contextmanager
def timed_task(task_name: str, **context_info: Any) -> Any:
    pass


def my_func(my_timed_task: Callable[..., AbstractContextManager[Any]]) -> Any:
    with my_timed_task("my_func_task", foo="bar"):
        pass

Also makes mypy happy. However there is no information about what arguments your my_timed_task factory accepts, so your IDE will not complain about doing something like this:

def my_func(my_timed_task: Callable[..., AbstractContextManager[Any]]) -> Any:
    with my_timed_task(1, 3, "a", True):
        pass

Which one of those options you choose is obviously up to you and depends on how precise you want to be.

Hope this helps.

user_197707950
Related