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