Support function decorators excellently

[sixolet]

Decorators currently stress mypy's support for functional programming to the point that many decorators are impossible to type. I'm intending this issue as a project plan and exploration of the issue of decorators and how to type them. It's a collection point for problems or incompletenesses that keep decorators from being fully supported, plans for solutions to those problems, and fully-implemented solutions to those problems.

Decorators that can only decorate a function of a fixed signature

You can define your decorator's signature explicitly, and it's completely supported:

from typing import Callable, Tuple

def with_strlen(f: Callable[[int], str]) -> Callable[[int], Tuple[str, int]]:
    def ret(__i: int) -> Tuple[str, int]:
        r = f(__i)
        return r, len(r)
    return ret

@with_strlen
def lol(x: int) -> str:
    return "lol"*x

reveal_type(lol)  # E: Revealed type is 'def (builtins.int) -> Tuple[builtins.str, builtins.int]

Notes:

• You can even use type variables to get some flexibility in argument and return types, but this falls over as soon as you don't know the exact number of arguments to expect.
• Your decorated function's arguments can't be called by name.
  • Better callable: Callable[[Arg('x', int), VarArg(str)], int] now a thing you can do #2607 enables that

Decorators that do not change the signature of the function

Nearly fully supported. Here's how you do it:

from typing import TypeVar, Callable, cast

T = TypeVar('T')

def print_callcount(f: T) -> T:
    x = 0
    def ret(*args, **kwargs):
        nonlocal x
        x += 1
        print("%d calls so far" % x)
        return f(*args, **kwargs)

    return cast(T, ret)

@print_callcount
def lol(x: int) -> str:
    return "lol"*x

reveal_type(lol)  # E: Revealed type is 'def (x: builtins.int) -> builtins.str'

Notes:

• Mypy trusts that you can call f. You can set a bound on T to be Callable, but you don't need to for it to typecheck.
• This business doesn't typecheck without the cast. That's a symptom of the fact that mypy doesn't understand the function nature of the argument to print_callcount, and there's no way to declare that argument as a Callable explicitly without losing the argument type of lol later. We'd like to minimize the places where casts are required. Doing so requires something along the lines of "variadic argument variables", discussed below.

Decorators that take arguments ("second-order"?)

Plenty of decorators "take arguments" by actually being functions that return a decorator. For example, we'd like to be able to do this:

from typing import Any, TypeVar, Callable, cast

T = TypeVar('T')

def callback_callcount(cb: Callable[[int], None]) -> Callable[[T], T]:
    def outer(f: T) -> T:
        x = 0
        def inner(*args, **kwargs):
            nonlocal x
            x += 1
            cb(x)
            return f(*args, **kwargs)
        return cast(T, inner)
    return outer

def print_int(x: int) -> None:
    print(x)

@callback_callcount(print_int)
def lol(x: int) -> str:
    return "lol"*x

reveal_type(lol)  # E: Revealed type is 'def (x: builtins.int) -> builtins.str'

Notes:

• This does not typecheck yet -- errors on calling the decorator, and lol ends up typed as None
• Relevant issue: Function returning a generic function #1551
  • TypeVar to represent a Callable's arguments #3028 fixes this.
• Still has a non-ideal cast...

Mess with the return type or with arguments

For an arbitrary function you can't do this at all yet -- there isn't even a syntax. Here's me making up some syntax for it.

Messing with the return type

from typing import Any, Dict, Callable

from mypy_extensions import SomeArguments

def reprify(f: Callable[[SomeArguments], Any]) -> Callable[[SomeArguments], str]:
    def ret(*args: SomeArguments.positional, **kwargs: SomeArguments.keyword):
        return repr(f(*args, **kwargs))
    return ret

@reprify
def lol(x: int) -> Dict[str, int]:
    return {"lol": x}

reveal_type(lol)  # E: Revealed type is 'def (x: builtins.int) -> builtins.str'

Messing with the arguments

from typing import Any, Callable, TypeVar

from mypy_extensions import SomeArguments

R = TypeVar('R')

def supply_zero(f: Callable[[int, SomeArguments], R]) -> Callable[[SomeArguments], R]:
    def ret(*args: SomeArguments.positional, **kwargs: SomeArguments.keyword):
        return f(0, *args, **kwargs)
    return ret

@supply_zero
def lol(x: int, y: str) -> str:
    return "%d and %s" % (x, y)

reveal_type(lol)  # E: Revealed type is 'def (y: builtins.str) -> builtins.str'

The syntax here is fungible, but we would need a way to do approximately this thing -- capture the types and kinds of all a function's arguments in some kind of variation on a type variable.

Relevant issues and discussions:

• Making a decorator which preserves function signature #1927
• TypeVar to represent a Callable's arguments #3028
• Make Callable more flexible typing#239
• Allow variadic generics typing#193

Variadic type variables alone (python/typing#193) get you some of the way there, but lose all keyword arguments of the decorated function.

Things to do:

• Implement variadic type variables (fill in PR when I have it)
• Write up detailed proposal for semantics of argument variables
  • ... and how they interact with *args and **kwargs
  • ... and their relationship to variadic type variables and the expand operation
  • ... and the semantics of an easy-to-use SomeArguments-style alias, so nobody has to actually engage with the details of the above when writing normal decorators.
• Come to some kind of mypy-community consensus or near-consensus on that proposal. It'll be in mypy_extensions not typing at first -- this can be fodder for the future of PEP484, but while we're playing in such experimental land, not yet.
• PR to implment the SomeArguments thing.

[gvanrossum]

Great to see this may be really happening! The Twitter discussion gave me a new suggestion for what to name "second-order decorators" -- people seemed to like "decorator factory" best.

[elazarg]

"decoratory" :)

[JukkaL]

Here are some thoughts about decorators that tweak arguments/return types. This proposal combines features from @sixolet's proposal and other sources, such as #1927 and python/typing#193.

Argspec type variables

We'd add a new kinds of type variable: argspec. An argspec type variable represents an arbitrary sequence of callable argument types/kinds, such as DefaultArg(int, 'b') using the syntax recently implemented by @sixolet. (We could also generalize this to variadic type variables that represent simple sequences of types, but that's out of scope for this issue.)

This is how we'd define one:

from typing import TypeVar

Args = TypeVar('Args', argspec=True)

For example, we could use this in a stub for contextlib to define the signature of contextmanager:

RT = TypeVar('RT')  # Regular type variable

def contextmanager(
        func: Callable[Args, Iterator[RT]]) -> Callable[Args, GeneratorContextManager[RT]]: ...

Decorator implementations

Outside a stub, we hit a problem: it's going to be tricky to type check implementations of functions that use an argspec type variable in their signatures. We can sidestep this issue by not supporting this at all -- instead we'd require a separate external signature declaration and an implementation for such functions, and the implementation can't use argspecs. This is analogous to how overloaded functions with implementations work. Sketch of how to implement contextmanager:

from typing import declared_type

...

@declared_type
def contextmanager(
    func: Callable[Args, Iterator[RT]]) -> Callable[Args, GeneratorContextManager[RT]]: ...

def contextmanager(
        func: Callable[..., Iterator[RT]]) -> Callable[..., GeneratorContextManager[RT]]:
    <implementation of context manager>
    return <something>

The name declared_type is analogous to decorated_type (see #3291), but I'm not convinced that this the best name we can think of. Maybe we should even combine declared_type and decorated_type.

__call__

Now what about a decorator that is not a callable but a more general object? We could still generalize this approach by making it possible to use an argspec to define the type of __call__ (here I assume a stub file):

class dec(Generic[Args, RT]):
    def __init__(self, fn: Callable[Args, RT]) -> None: ...
    def __call__(self, *args: Args, **kwargs: Args) -> RT: ...

Note that Args can be used for both *args and **kwargs -- actually, it must always be used in such a way if used outside a Callable type, as only *args: Args doesn't really make sense.

Expand

We can support an Expand[...] operator that lets us tweak argument lists. For example, here we add an extra int argument to the beginning of an argument list:

def dec(x: Callable[Args, RT]) -> Callable[[int, Expand[Args]], RT]: ...

The relationship between Args and Expand[Args] is like the relationship between alist and *alist.

[JukkaL]

Oh and the updated TypeVar and declared_type would initially live in mypy_extensions.

[ilevkivskyi]

I have two clarifying questions/proposals:

• Maybe we can allow a simpler version of Expand for Python 3.5 and later:

  def dec(x: Callable[Args, RT]) -> Callable[[int, *Args], RT]: ...

• How should be defined decorators that reduce number of variables? Just in an opposite way?

  def dec(x: Callable[[int, *Args], RT]) -> Callable[Args, RT]: ...

[rkr-at-dbx]

Outside a stub, we hit a problem: it's going to be tricky to type check implementations of functions that use an argspec type variable in their signatures.

How so? I'm aware that it's difficult for general variadic functions (as mentioned in python/typing#193), but I'd expect that most decorators will just treat values of Args type as opaque objects to be passed on to their wrapped function.

[sixolet]

@JukkaL I have been pondering this problem for the last two weeks, and haven't started an implementation of any solution because I hadn't solved it yet in my head. I like a whole lot about your suggestions, and I think they get us a lot further.

Some wiggly bits I have:

• I don't think argspec is a kind of type variable. I think it's different enough it should be its own thing.

• When you have this code block:

from typing import declared_type

@declared_type
def contextmanager(
    func: Callable[Args, Iterator[RT]]) -> Callable[Args, GeneratorContextManager[RT]]: ...

def contextmanager(
        func: Callable[..., Iterator[RT]]) -> Callable[..., GeneratorContextManager[RT]]:
    <implementation of context manager>
    return <something>

Can we consider this instead:

contextmanager_sig = Callable[
    [Callable[Args, Iterator[RT]], 
    Callable[Args, GeneratorContextManager[RT]]
]

@declared_type(contextmanager_sig)
def contextmanager(
        func: Callable[..., Iterator[RT]]) -> Callable[..., GeneratorContextManager[RT]]:
    <implementation of context manager>
    return <something>

This uses a more direct method than redefinition to declare the type of the function aside from its implementation. It also neatly sidesteps the problem of how to define an argspec in the def syntax at all, for some weird signatures:

@declared_type(Callable[[Callable[Args, R], Expand[Args]], R)
def apply(f: Callable[..., R], *args, **kwargs) -> R: ...

... I don't think these weird signatures are important enough to specifically jump through hoops on their own. Should do this form if we feel like it's cleaner.

The following is speculation and musing and I'm not particularly strongly attached to it.

I have some kind of sneaking suspicion that what we're calling "argspecs" here and variadic type variables are more closely related than (at least I) suspected. I'm looking for a way to relate these two concepts -- for example, a way to cleanly and pleasantly write the signature of map.

Maybe something like:

Args = ArgSpec('ArgSpec', keyword=False) # matches only positional arguments
R = TypeVar('R')

@declared_type(Callable[[Callable[Args, R], Expand[Iterable[Args], Args]], Iterable[R]])
def map(f: Callable[..., R], *args) -> Iterable[R]: ...

(Here Expand takes an second argument, specifically the argspec/variadic type var to expand. This turns it into some kind of comprehension.)

• Whichever of these forms we pick, the form for declaring the type of a decorated function should be exactly the same -- they shouldn't be two concepts.

[JukkaL]

@sixolet

I don't think argspec is a kind of type variable. I think it's different enough it should be its own thing.

That's fair. Anyway, we don't need to decide this very early, since this is only a matter of syntax.

This uses a more direct method than redefinition to declare the type of the function aside from its implementation.

I have no strong opinion either way right now, but I have a feeling that the flexible callable syntax can make more complex signatures hard to read. Again, this is mostly a matter of syntax and we can bikeshed it later after we've agreed on basic principles. And I agree that making this consistent with how we declare the decorated signature of a function would be nice.

I have some kind of sneaking suspicion that what we're calling "argspecs" here and variadic type variables are more closely related than (at least I) suspected.

Yes, I think that a variadic type variables would mostly behave like an argspec, but it wouldn't have any argument kinds (just a sequence of types). I left this out from my proposal because it's not very directly related to the current issue. A variadic type variable would also allow things like Tuple[Expand[X]], which argspecs wouldn't support.

@rkr-at-dbx

How so? I'm aware that it's difficult for general variadic functions (as mentioned in python/typing#193), but I'd expect that most decorators will just treat values of Args type as opaque objects to be passed on to their wrapped function.

Yes, most decorators probably are trivial in that respect. However, any non-trivial decorators could be very tricky to type check. I'm worried that users have a significant number of those (though still a minority of all decorators) and if we don't type check them properly it will be a never-ending stream of bug reports and ideas about handling various edge cases. Also, as these are otherwise quite similar to variadic type variables, it might feel a little odd if we could type check one but not the other.

@ilevkivskyi

Maybe we can allow a simpler version of Expand for Python 3.5 and later:

Unfortunately that doesn't work with variadic type variables for things like Tuple[Expand[X]], and I'd rather not have two different syntax variants for expanding type variables.

How should be defined decorators that reduce number of variables? Just in an opposite way?

Yes, though changing the number of arguments might be something we implement only a bit later as the feature would be quite useful even without it.