Augmenting third-party packages

[flying-sheep]

Motivation

Some framework packages (like pytest, polars, xarray, …) have APIs that allow plugin packages to define attributes on their classes/singletons using some registration function, e.g.:

import pytest

def pytest_configure(config: pytest.Config) -> None:
    config.addinivalue_line("markers", "mychoice(select, skip): choose which stuff to test")

@pytest.mark.mychoice(x=1)  # I want to make this into a type error by defining the signature somewhere
def test_thing(): ...

These plugins should have a way to specify the type of the new attribute. In our example, pytest itself has this definition for the type of the pytest.mark object:

@final
class MarkGenerator:
    if TYPE_CHECKING:
        skip: _SkipMarkDecorator
        skipif: _SkipifMarkDecorator
        xfail: _XfailMarkDecorator
        parametrize: _ParametrizeMarkDecorator
        usefixtures: _UsefixturesMarkDecorator
        filterwarnings: _FilterwarningsMarkDecorator

    # untyped marks:
    def __getattr__(self, name: str) -> MarkDecorator: ...

which allows its own defined marks to be typed:

@pytest.mark.skipif(x=None)  # this *is* a type error
def test_something(): ...

A plugin needs to have a way to add a new typed attribute to MarkGenerator.

Design considerations

The current way of shipping types has no good way of having potentially multiple stubs that can be merged into one: even if it’s possible to ship pytest/__init__.pyi in one plugin and have it merged with the actual pytest package, only one plugin could do that, and there would be no indication that this .pyi is intended to be an augmentation instead of a replacement for all of pytest’s types.

So we’d need a new way to locate augmentation stubs, I think.

As for how these stubs look like, I think typing.Protocol could do a good job:
A Protocol in an augmentation stub could be interpreted as an augmentation protocol, e.g.

$PYTHONPATH/my-pytest-plugin/typeshedding-location-for-augments/pytest/__init__.pyi or
$PYTHONPATH/typeshedding-location-for-augments/my-pytest-plugin/pytest/__init__.pyi

from typing import Protocol
import pytest

class MarkGenerator(Protocol):
    mychoice: _MyChoiceMarkDecorator

class _MyChoiceMarkDecorator(pytest.MarkDecorator):
    def __call__(  # type: ignore[override]
            self,
            select: list[str] = ...,
            skip: list[str] = ...,
        ) -> MarkDecorator: ...

Prior art

• TypeScript has this feature, explained here: https://www.typescriptlang.org/docs/handbook/declaration-merging.html#module-augmentation
• PEP 484 had a location for distributing stubs outside of the actual package, we could bring it back for augments

[erictraut]

The TypeScript feature you cite (implicit module augmentation) has posed big problems for TypeScript performance and usability. If you were to ask the folks who created this feature for TypeScript, they'd likely tell you that it's a feature they regret — or would design differently if they had a chance to do so. It requires that the TypeScript compiler open, read, parse, and bind every file in a project before it can perform any type analysis on one file. This is problematic when opening a large project in an editor with language server support. It can take many seconds or even minutes to get syntax highlighting, completion suggestions, etc. when opening such a project. Implicit augmentation also creates usability problems because it can have side effects that are difficult to debug and report. For all of those reasons, I would discourage going down the path of implicit augmentation.

Let's consider options that are more explicit and therefore avoid the downsides I discuss above.

The first tool that comes to mind is intersection types. This facility, which has been discussed and explored extensively in the typing forums, allows a type to take on the properties of two or more existing types. This is a good way to add new attributes to an existing class, for example. And there are many use cases that have already been identified for intersection types.

Merging of stubs is challenging, even if the locations of these stubs are specified explicitly. Such merging would introduce the potential for errors that would be difficult to explain to the user. This is where the usability issues crop up in TypeScript as well. An alternative solution that avoids this issue is to rely on the developer to create a local module that both registers plugins and defines new types using intersections.

# plugins.py
import pytest
import plugin1
import plugin2

plugin1.register()
plugin2.register()

# The ExtendedDecorator type can be used elsewhere in the project
# where a pytest MarkDecorator would normally be used.
# Note: I'm using the `&` operator here to denote "intersection".
type ExtendedDecorator = pytest.MarkDecorator & plugin1.DecoratorExtension & plugin2.DecoratorExtension

Do you think that something along those lines would meet your requirements?

[flying-sheep]

Let's consider options that are more explicit and therefore avoid the downsides I discuss above.

for sure, that was my suggestion as well: augments as an explict feature that can’t happen in literally any stub anywhere.

An alternative solution that avoids this issue is to rely on the developer to create a local module that both registers plugins and defines new types using intersections.

that’s not a solution, because the type of pytest.mark stays the same. An actual solution would have to be able to actually modify that type. So nothing along those lines can meet my requirements.

Merging of stubs is challenging, even if the locations of these stubs are specified explicitly. Such merging would introduce the potential for errors that would be difficult to explain to the user.

Yeah, but: type theory is challenging, doing it right is challenging, making things ergonomic is challenging, but we have to do all of it.

---

So with that out of the way, let‘s think about how an actual solution with optimal performance would look like!

My suggestion above still stands: allow each package to augment each other package, with augments being an explicitly incomplete kind of type stub.

• performance is easy to keep in check: each package has only one chance to augment another (see paths above: $PYTHONPATH/my-package/typeshedding-location-for-augments/augmented-package/…). the checker can cache which packages have augments at all, and can quickly check which package augments which.
• debug ergonomics don’t sound that hard either: type checkers could grow logging or command line switches that list the locations where types for a certain package come from.