How to define a ContextManager Protocol

Question

I am trying to use type hinting to specify the API to follow when implementing a connector class (to a broker, in this case).

I want to specify that such class(es) should be context manager(s)

How do I do that?

Let me reword it more clearly: how can I define the Broker class so that it indicates that its concrete implementations, e.g. the Rabbit class, must be context managers?

Is there a practical way? Do I have to specify __enter__ and __exit__ and just inherit from Protocol?

Is it enough to inherit from ContextManager?

By the way, should I use @runtime or @runtime_checkable? (My VScode linter seems to have problems finding those in typing. I am using python 3 7.5)

I know how to do it with ABC's, but I would like to learn how to do it with protocol definitions (which I have used fine already, but they weren't context managers).

I cannot make out how to use the ContextManager type. So far I haven't been able to find good examples from the official docs.

At present I came up with

from typing import Protocol, ContextManager, runtime, Dict, List


@runtime
class Broker(ContextManager):
    """
    Basic interface to a broker.
    It must be a context manager
    """

    def publish(self, data: str) -> None:
        """
        Publish data to the topic/queue
        """
        ...

    def subscribe(self) -> None:
        """
        Subscribe to the topic/queue passed to constructor
        """
        ...

    def read(self) -> str:
        """
        Read data from the topic/queue
        """
        ...

and the implementation is

@implements(Broker)
class Rabbit:
    def __init__(self,
            url: str,
            queue: str = 'default'):
        """
        url: where to connect, i.e. where the broker is
        queue: the topic queue, one only
        """
        # self.url = url
        self.queue = queue
        self.params = pika.URLParameters(url)
        self.params.socket_timeout = 5

    def __enter__(self):
        self.connection = pika.BlockingConnection(self.params) # Connect to CloudAMQP
        self.channel = self.connection.channel() # start a channel
        self.channel.queue_declare(queue=self.queue) # Declare a queue
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.connection.close()

    def publish(self, data: str):
        pass  # TBD

    def subscribe(self):
        pass  # TBD

    def read(self):
        pass  # TBD

Note: the implements decorator works fine (it comes form a previous project), it checks the class is a subclass of the given protocol

Answer 1

Short answer -- your Rabbit implementation is actually fine as-is. Just add some type hints to indicate that __enter__ returns an instance of itself and that __exit__ returns None. The types of the __exit__ params don't actually matter too much.

---

Longer answer:

Whenever I'm not sure what exactly some type is/what some protocol is, it's often helpful to check TypeShed, the collection of type hints for the standard library (and a few 3rd party libraries).

For example, here is the definition of typing.ContextManager. I've copied it below here:

from types import TracebackType

# ...snip...

_T_co = TypeVar('_T_co', covariant=True)  # Any type covariant containers.

# ...snip...

@runtime_checkable
class ContextManager(Protocol[_T_co]):
    def __enter__(self) -> _T_co: ...
    def __exit__(self, __exc_type: Optional[Type[BaseException]],
                 __exc_value: Optional[BaseException],
                 __traceback: Optional[TracebackType]) -> Optional[bool]: ...

From reading this, we know a few things:

1. This type is a Protocol, which means any type that happens to implement __enter__ and __exit__ following the given signatures above would be a valid subtype of typing.ContextManager without needing to explicitly inherit it.

2. This type is runtime checkable, which means doing isinstance(my_manager, ContextManager) also works, if you want to do that for whatever reason.

3. The parameter names of __exit__ are all prefixed with two underscores. This is a convention type checkers use to indicate those arguments are positional only: using keyword arguments on __exit__ won't type check. Practically speaking, that means you can name your own __exit__ parameters whatever you'd like while still being in compliance with the protocol.

So, putting it together, here is the smallest possible implementation of a ContextManager that still type checks:

from typing import ContextManager, Type, Generic, TypeVar

class MyManager:
    def __enter__(self) -> str:
        return "hello"

    def __exit__(self, *args: object) -> None:
        return None

def foo(manager: ContextManager[str]) -> None:
    with manager as x:
        print(x)        # Prints "hello"
        reveal_type(x)  # Revealed type is 'str'

# Type checks!
foo(MyManager())



def bar(manager: ContextManager[int]) -> None: ...

# Does not type check, since MyManager's `__enter__` doesn't return an int
bar(MyManager())

One nice little trick is that we can actually get away with a pretty lazy __exit__ signature, if we're not actually planning on using the params. After all, if __exit__ will accept basically anything, there's no type safety issue.

(More formally, PEP 484 compliant type checkers will respect that functions are contravariant with respect to their parameter types).

But of course, you can specify the full types if you want. For example, to take your Rabbit implementation:

# So I don't have to use string forward references
from __future__ import annotations
from typing import Optional, Type
from types import TracebackType

# ...snip...

@implements(Broker)
class Rabbit:
    def __init__(self,
            url: str,
            queue: str = 'default'):
        """
        url: where to connect, i.e. where the broker is
        queue: the topic queue, one only
        """
        # self.url = url
        self.queue = queue
        self.params = pika.URLParameters(url)
        self.params.socket_timeout = 5

    def __enter__(self) -> Rabbit:
        self.connection = pika.BlockingConnection(params) # Connect to CloudAMQP
        self.channel = self.connection.channel() # start a channel
        self.channel.queue_declare(queue=self.queue) # Declare a queue
        return self

    def __exit__(self,
                 exc_type: Optional[Type[BaseException]],
                 exc_value: Optional[BaseException],
                 traceback: Optional[TracebackType],
                 ) -> Optional[bool]:
        self.connection.close()

    def publish(self, data: str):
        pass  # TBD

    def subscribe(self):
        pass  # TBD

    def read(self):
        pass  # TBD

---

To answer the new edited-in questions:

How can I define the Broker class so that it indicates that its concrete implementations, e.g. the Rabbit class, must be context managers?

Is there a practical way? Do I have to specify enter and exit and just inherit from Protocol?

Is it enough to inherit from ContextManager?

There are two ways:

1. Redefine the __enter__ and __exit__ functions, copying the original definitions from ContextManager.
2. Make Broker subclass both ContextManager and Protocol.

If you subclass only ContextManager, all you are doing is making Broker just inherit whatever methods happen to have a default implementation in ContextManager, more or less.

PEP 544: Protocols and structural typing goes into more details about this. The mypy docs on Protocols have a more user-friendly version of this. For example, see the section on defining subprotocols and subclassing protocols.

By the way, should I use @runtime or @runtime_checkable? (My VScode linter seems to have problems finding those in typing. I am using python 3 7.5)

It should be runtime_checkable.

That said, both Protocol and runtime_checkable were actually added to Python in version 3.8, which is probably why your linter is unhappy.

If you want to use both in older versions of Python, you'll need to pip install typing-extensions, the official backport for typing types.

Once this is installed, you can do from typing_extensions import Protocol, runtime_checkable.