Make `@overload` less verbose

Overloading can become extremely verbose very fast, because all arguments have to be repeated every time. However, often only the type hints on a few arguments change.

It would be nice to be able to abbreviate overloads by skipping arguments, and deferring their type hints. Consider this example from pandas-stubs.

Series.reset_index type hint
    @overload
    def reset_index(
        self,
        level: Sequence[Level] = ...,
        *,
        drop: Literal[False] = ...,
        name: Level = ...,
        inplace: Literal[False] = ...,
        allow_duplicates: bool = ...,
    ) -> DataFrame: ...
    @overload
    def reset_index(
        self,
        level: Sequence[Level] = ...,
        *,
        drop: Literal[True],
        name: Level = ...,
        inplace: Literal[False] = ...,
        allow_duplicates: bool = ...,
    ) -> Series[S1]: ...
    @overload
    def reset_index(
        self,
        level: Sequence[Level] = ...,
        *,
        drop: bool = ...,
        name: Level = ...,
        inplace: Literal[True],
        allow_duplicates: bool = ...,
    ) -> None: ...
    @overload
    def reset_index(
        self,
        level: Level | None = ...,
        *,
        drop: Literal[False] = ...,
        name: Level = ...,
        inplace: Literal[False] = ...,
        allow_duplicates: bool = ...,
    ) -> DataFrame: ...
    @overload
    def reset_index(
        self,
        level: Level | None = ...,
        *,
        drop: Literal[True],
        name: Level = ...,
        inplace: Literal[False] = ...,
        allow_duplicates: bool = ...,
    ) -> Series[S1]: ...
    @overload
    def reset_index(
        self,
        level: Level | None = ...,
        *,
        drop: bool = ...,
        name: Level = ...,
        inplace: Literal[True],
        allow_duplicates: bool = ...,
    ) -> None: ...

The many overloads with many arguments make it hard to see at a glance what is going on. Compare this with the following, with some simplifications:

@overload
def reset_index(
    ...,
    *,
    inplace: Literal[True],
    ...,
) -> None: ...
@overload
def reset_index(
    ..., 
    *,
    drop: Literal[True],
    ...,
) -> Series[S1]: ...
@overload
def reset_index(
    ...,
    *,
    drop: Literal[False],
    ...,
) -> DataFrame: ...
@overload
def reset_index(
    self,
    level: Sequence[Level] | Level | None = ...,
    *,
    drop: bool = ...,
    name: Level = ...,
    inplace: bool,
    allow_duplicates: bool = ...,
) -> None | Series[S1] | DataFrame: ...

Which immediately makes it clear which arguments give which return. How this should work is when a ... is encountered in an @overload, then the type-hints of missing arguments are deferred to later overloads. Obviously this is just a very rough idea, but I wonder if other people feel the same.

Another benefit could be with refactoring. If an extra argument is added, we only have to modify the last overload instead of adding it to every existing overload.

6 Likes

While I do like the terseness of this idea, unfortunately it will not be possible without changes to the syntax, which makes it harder to justify, especially since ... outside of a function overload would be nonsensical but it would still have to be valid syntax otherwise. So I don’t think this makes sense unless overloads become their own AST node with their own grammar rules.

Maybe we can come up with a way that’s syntactically valid, non-ambiguous and still relatively terse.

Maybe something like *_: auto and **_: auto meaning the elided arguments would be determined by the following overloads. The issue with that approach however is that we can’t easily and non-ambiguously elide positional arguments, maybe we’d need something like _: ellide(n) to skip n positional arguments.

There’s an existing thread where I also shared some of my thoughts about how to at least reduce the number of overloads you have to write due to the complications introduced through arguments that can either be positional or keyword: Some improvements to function overloads

Python 3.3.7+ (default, Jun 21 2022, 18:59:46) 
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> ...
Ellipsis

Problem solved. :slight_smile:

I think the point was that it would have to be valid syntax as part of the grammar for a function header, in the place where currently an identifier is expected.

>>> def test(...): pass
  File "<stdin>", line 1
    def test(...): pass
             ^
SyntaxError: invalid syntax

It occurs to me that if there were some kind of built-in type that represented the arguments to a function, along with a way to collect them (regardless of whether the parameters use *args/**kwargs, then one could use match-case for the dispatch instead of @overload. But maybe that isn’t particularly helpful. (It seems like that way would involve putting all the implementations inside the same function body, or else writing even more boilerplate.)

1 Like

Yes, sorry if I wasn’t being clear enough. I didn’t mean an ellipsis on its own would be nonsensical, I meant this proposed new syntax would be nonsensical in a function that has no overloads, so it would be strange to change the grammar for something that only does something in function signatures with the overload decorator, unless you also introduced a new syntax to specify overloads, so the grammar change could be applied only to overloads and not regular function signatures as well.

def foo(...): ...

This would need to be valid syntax, but what does that mean? It could make sense as a shorthand for:

def foo(*_, **_): ...

i.e. accepting any arguments but ignoring them inside the function body, but it becomes less obvious once you start mixing in other function signature syntax like / * and **. Also that would be kind of contradictory with the proposed meaning in overloads. In overloads it’s supposed to mean “Infer the ellided arguments from the other overloads”, which I’m not even sure is actually possible to do completely non-ambiguously, unless you specify how many arguments were ellided, since some overloads may omit a subset of the arguments entirely or choose different names for the same positional only argument to provide better documentation.

2 Likes

I believe

def test(...):
    ...

could be valid syntax for parser, but calling such a function should fail. In practice the signature object of the parsed function could contain a parameter with a kind having some new special value like PLACEHOLDER. Then calling such a function would fail with a TypeError having a message like:

Cannot call a function with a placeholder in its signature

This would make the desired @overflow syntax valid, but could be used also for other purposes.

2 Likes

The ... syntax could have other uses as well. For example, currently it seems not possible to use the built-in abstract base classes and typing module to specify a Protocol / ABC that specifies that subclasses must implement a function foo, without putting any constraints on the signature of foo.

Consequently, libraries that need this like pytorch use workarounds like

class Module:
    forward: Callable[..., Any] = _forward_unimplemented

with the Ellipsis-syntax, something like the following should be possible

class Module
    @abstractmethod
    def forward(...) -> Any: ...

Which type checkers can interpret as weakest upper bound of all callables type[forward] = Callable[..., Any]. Similarly, it could enable specification of partial signatures, which is currently mostly unsupported.

2 Likes

+1 for making overload less verbose. The current syntax is hard to parse and maintain when there’s multiple arguments.

Is it necessary for this last overload to repeat the return type of None | Series[S1] | DataFrame when the prior overloads already specified it collectively? How about allowing the return type to be ... here to keep the code DRY?

By the same token, it can be argued that the type of drop can also be left as ... if the prior overloads have sufficiently specified all the possibilities.

I agree. Since ... can’t reasonably be used as part of the header of a regular function definition it may justify its own syntax.

How about a dedicated decl/declaration statement, with no more need for the boilerplate of a body : ...:

decl reset_index(..., *, drop: Literal[True], ...) -> Series[S1]
decl reset_index(..., *, drop: Literal[False], ...) -> DataFrame
decl reset_index(self, level: Level, *, drop: ...) -> ...

And maybe we can simply allow a declaration to not type a parameter or the return value at all to indicate that the type is sufficiently specified in the other declarations, so the last declaration above can be written as:

decl reset_index(self, level: Level, *, drop)

There are (too) many examples of “extreme overloading” in scipy-stubs, for instance, the scipy.sparse._base._spbase.__mul__ signature looks like this:

@overload  # Self[-Bool], /, other: scalar-like +Bool
def __mul__(self, /, other: bool | _ToBool) -> Self: ...
@overload  # Self[-Int], /, other: scalar-like +Int
def __mul__(self: _SpFromIntT, /, other: onp.ToInt) -> _SpFromIntT: ...
@overload  # Self[-Float], /, other: scalar-like +Float
def __mul__(self: _SpFromFloatT, /, other: onp.ToFloat) -> _SpFromFloatT: ...
@overload  # Self[-Complex], /, other: scalar-like +Complex
def __mul__(self: _SpFromComplexT, /, other: onp.ToComplex) -> _SpFromComplexT: ...
@overload  # sparray[-Bool], /, other: sparse +Bool
def __mul__(self: _SpArray, /, other: _spbase[_ToBool | _SCT_co]) -> _SpArrayOut[_SCT_co]: ...
@overload  # sparray[-Bool], /, other: array-like +Bool
def __mul__(self: _SpArray, /, other: _To2DLike[bool, _ToBool]) -> coo_array[_SCT_co, _ShapeT_co]: ...
@overload  # sparray[-Int], /, other: sparse +Int
def __mul__(self: _SpArray[_FromInt], /, other: _spbase[_ToInt8 | _SCT_co]) -> _SpArrayOut[_SCT_co]: ...
@overload  # sparray[-Int], /, other: array-like +Int
def __mul__(self: _SpArray[_FromInt], /, other: _To2DLike[bool, _ToInt8]) -> coo_array[_SCT_co, _ShapeT_co]: ...
@overload  # sparray[-Float], /, other: sparse +Float
def __mul__(self: _SpArray[_FromFloat], /, other: _spbase[_ToFloat32 | _SCT_co]) -> _SpArrayOut[_SCT_co]: ...
@overload  # sparray[-Float], /, other: array-like +Float
def __mul__(self: _SpArray[_FromFloat], /, other: _To2DLike[int, _ToFloat32]) -> coo_array[_SCT_co, _ShapeT_co]: ...
@overload  # sparray[-Complex], /, other: sparse +Complex
def __mul__(self: _SpArray[_FromComplex], /, other: _spbase[_ToComplex64 | _SCT_co]) -> _SpArrayOut[_SCT_co]: ...
@overload  # sparray[-Complex], /, other: array-like +Complex
def __mul__(self: _SpArray[_FromComplex], /, other: _To2DLike[int, _ToComplex64]) -> coo_array[_SCT_co, _ShapeT_co]: ...
@overload  # spmatrix, /, other: spmatrix
def __mul__(self: _SpMatrixT, /, other: _SpMatrixT) -> _SpMatrixT: ...
@overload  # spmatrix[-Bool], /, other: sparse +Bool
def __mul__(self: spmatrix, /, other: _spbase[_ToBool]) -> _SpMatrixOut[_SCT_co]: ...
@overload  # spmatrix[-Bool], /, other: array-like +Bool
def __mul__(self: spmatrix, /, other: _To2D[bool, _ToBool]) -> onp.Array2D[_SCT_co]: ...
@overload  # spmatrix[-Int], /, other: sparse +Int
def __mul__(self: spmatrix[_FromInt], /, other: _spbase[_ToInt8]) -> _SpMatrixOut[_SCT_co]: ...
@overload  # spmatrix[-Int], /, other: array-like +Int
def __mul__(self: spmatrix[_FromInt], /, other: _To2D[bool, _ToInt8]) -> onp.Array2D[_SCT_co]: ...
@overload  # spmatrix[-Float], /, other: sparse +Float
def __mul__(self: spmatrix[_FromFloat], /, other: _spbase[_ToFloat32 | _SCT_co]) -> _SpMatrixOut[_SCT_co]: ...
@overload  # spmatrix[-Float], /, other: array-like +Float
def __mul__(self: spmatrix[_FromFloat], /, other: _To2D[int, _ToFloat32]) -> onp.Array2D[_SCT_co]: ...
@overload  # spmatrix[-Complex], /, other: sparse +Complex
def __mul__(self: spmatrix[_FromComplex], /, other: _spbase[_ToComplex64 | _SCT_co]) -> _SpMatrixOut[_SCT_co]: ...
@overload  # spmatrix[-Complex], /, other: array-like +Complex
def __mul__(self: spmatrix[_FromComplex], /, other: _To2D[float, _ToComplex64]) -> onp.Array2D[_SCT_co]: ...
@overload  # spmatrix[+Bool], /, other: scalar- or matrix-like ~Int
def __mul__(self: spmatrix[_ToBool], /, other: _SparseLike[opt.JustInt, Int]) -> spmatrix[Int]: ...
@overload  # spmatrix[+Bool], /, other: array-like ~Int
def __mul__(self: spmatrix[_ToBool], /, other: _To2D[opt.JustInt, Int]) -> onp.Array2D[Int]: ...
@overload  # spmatrix[+Int], /, other: scalar- or matrix-like ~Float
def __mul__(self: spmatrix[_ToInt], /, other: _SparseLike[opt.JustFloat, Float]) -> spmatrix[Float]: ...
@overload  # spmatrix[+Int], /, other: array-like ~Float
def __mul__(self: spmatrix[_ToInt], /, other: _To2D[opt.JustFloat, Float]) -> onp.Array2D[Float]: ...
@overload  # spmatrix[+Float], /, other: scalar- or matrix-like ~Complex
def __mul__(self: spmatrix[_ToFloat], /, other: _SparseLike[opt.JustComplex, Complex]) -> spmatrix[Complex]: ...
@overload  # spmatrix[+Float], /, other: array-like ~Complex
def __mul__(self: spmatrix[_ToFloat], /, other: _To2D[opt.JustComplex, Complex]) -> onp.Array2D[Complex]: ...
@overload  # Self[+Bool], /, other: -Int
def __mul__(self: _spbase[_ToBool], /, other: _FromIntT) -> _spbase[_FromIntT, _ShapeT_co]: ...
@overload  # Self[+Int], /, other: -Float
def __mul__(self: _spbase[_ToInt], /, other: _FromFloatT) -> _spbase[_FromFloatT, _ShapeT_co]: ...
@overload  # Self[+Float], /, other: -Complex
def __mul__(self: _spbase[_ToFloat], /, other: _FromComplexT) -> _spbase[_FromComplexT, _ShapeT_co]: ...
@overload  # catch-all
def __mul__(self, /, other: _To2DLike[complex, Scalar] | _spbase) -> _spbase[Any, Any] | onp.Array[Any, Any]: ...

WIthout the comments this is practically impossible to work with. But even with them, it’s still extremely difficult to read.

But in this case, this proposed solution won’t help much, as there are only 2 parameters. The long-requested higher kinded type-parameters and intersection type would help, but I don’t think that having those features would help clean up your reset_index examplle from pandas-stubs.

An idea I’ve been playing with, that signifcantly help in both our cases, is something I can best describe as a “type-mapping”. It’s basically a generalization of a “typevar with constraints” such as AnyStr, which is a type-mapping of str -> str and bytes -> bytes.

So for reset_index, you could define a type-mapping that takes the inplace and drop types as input, and outputs the return type of reset_index. It could, for example, be written with some flashy scala-esque syntax (that shouldn’t be taken too seriously) like

from typing import Literal as L

case type ResetIndexResult[InplaceT: L[True], DropT: bool, VT] = None
case type ResetIndexResult[InplaceT: L[False], DropT: L[True], VT] = Series[VT]
case type ResetIndexResult[InplaceT: L[False], DropT: L[False], VT] = DataFrame

def reset_index[InplaceT: bool = L[False], DropT: bool = L[True]](
    self,
    level: Sequence[Level] | Level | None = ...,
    *,
    drop: DropT = ...,
    name: Level = ...,
    inplace: InplaceT = ..., 
    allow_duplicates: bool = ...,
) -> ResetIndexResult[InplaceT, DropT, S1]: ...

The scipy-stubs example has only 2 parameters, so I’m sure you can imagine how this type-mapping could also be applied there.

The advantage of this external definition, is that it’s reusable without depending on the specific parameter names or positions. In scipy-stubs, this could at the very least result in a 50% LOC reduction. This is also the case for the stubs of numpy, maybe even more so.

2 Likes