How to combine a custom protocol with the Callable protocol?

user_6117426 0
Viewed 2310

I have a decorator that takes a function and returns the same function with some added attributes:

import functools
from typing import *


def decorator(func: Callable) -> Callable:
    func.attr1 = "spam"
    func.attr2 = "eggs"
    return func

How do I type hint the return value of decorator? I want the type hint to convey two pieces of information:

  1. the return value is a Callable
  2. the return value has attributes attr1 and attr2

If I write a protocol,

class CallableWithAttrs(Protocol):
    attr1: str
    attr2: str

then I lose Callable. And apparently I can't make the protocol inherit from Callable;

class CallableWithAttrs(Callable, Protocol):
    attr1: str
    attr2: str

mypy says:

error: Invalid base class "Callable"

On the other hand, if I just use Callable, I lose the information about the added attributes.


This is perhaps even more complicated when introducing type variables, i.e. when the decorator must return the same type of callable as the given function func, as pointed out by MisterMiyagi in the comments.

import functools
from typing import *

C = TypeVar('C', bound=Callable)


def decorator(func: C) -> C:
    func.attr1 = "spam"
    func.attr2 = "eggs"
    return func

Now what do I do? I can't inherit from a type variable:

class CallableWithAttrs(C, Protocol):
    attr1: str
    attr2: str

error: Invalid base class "C"
python-typing type-hinting python
2 Answers

Since typing.Callable corresponds to collections.abc.Callable, you can just define a Protocol that implements __call__:

class CallableWithAttrs(Protocol):
    attr1: str
    attr2: str

    def __call__(self, *args, **kwargs): pass
user_37672390

One can parameterise a Protocol by a Callable:

from typing import Callable, TypeVar, Protocol

C = TypeVar('C', bound=Callable)  # placeholder for any Callable


class CallableObj(Protocol[C]):   # Protocol is parameterised by Callable C ...
    attr1: str
    attr2: str

    __call__: C                   # ... which defines the signature of the protocol

This creates an intersection of the Protocol itself with an arbitrary Callable.

A function that takes any callable C can thus return CallableObj[C], a callable of the same signature with the desired attributes:

def decorator(func: C) -> CallableObj[C]: ...

MyPy properly recognizes both the signature and attributes:

def dummy(arg: str) -> int: ...

reveal_type(decorator(dummy))           # CallableObj[def (arg: builtins.str) -> builtins.int]'
reveal_type(decorator(dummy)('Hello'))  # int
reveal_type(decorator(dummy).attr1)     # str
decorator(dummy)(b'Fail')  # error: Argument 1 to "dummy" has incompatible type "bytes"; expected "str"
decorator(dummy).attr3     # error: "CallableObj[Callable[[str], int]]" has no attribute "attr3"; maybe "attr2"?
user_53499160
Related