Let's assume I have a pydantic model, such as this Widget:

models.py

from pydantic import BaseModel


class Widget(BaseModel):
    name: str
    value: float

When writing tests (using pytest), which use this Widget, I frequently want to be able to create widgets on the fly, with some default values for its fields (which I do not want to set as default values in general, ie on the model, because they are only meant as default values for tests), and potentially some fields being set to certain values.

For this I currently have this clunky construct in my conftest.py file:

conftest.py

from typing import NotRequired, Protocol, TypedDict, Unpack

import pytest

from models import Widget


class WidgetFactoryKwargs(TypedDict):
    name: NotRequired[str]
    value: NotRequired[float]


class WidgetFactory(Protocol):
    def __call__(self, **kwargs: Unpack[WidgetFactoryKwargs]) -> Widget: ...


@pytest.fixture
def widget_factory() -> WidgetFactory:
    def _widget_factory(**kwargs: Unpack[WidgetFactoryKwargs]) -> Widget:
        defaults = WidgetFactoryKwargs(name="foo", value=42)
        kwargs = defaults | kwargs
        return Widget(**kwargs)

    return _widget_factory

This gives the type checker the ability to check if I am using the factory correctly in my tests and gives my IDE autocompletion powers:

test_widgets.py

from typing import assert_type

from conftest import WidgetFactory


def test_widget_creation(widget_factory: WidgetFactory) -> None:
    widget = widget_factory()
    assert_type(widget, Widget)        # during type checking
    assert isinstance(widget, Widget)  # during run time
    assert widget.name == "foo"
    assert widget.value == 42

    widget = widget_factory(name="foobar")
    assert widget.name == "foobar"
    assert widget.value == 42

    widget = widget_factory(value=1337)
    assert widget.name == "foo"
    assert widget.value == 1337

    widget = widget_factory(name="foobar", value=1337)
    assert widget.name == "foobar"
    assert widget.value == 1337

    widget = widget_factory(mode="maintenance")  # type checker error

(the actual tests are of course more involved and use the widget in some other way)

Question:

Is there a better way to achieve this type safety? Ideally, I could build the WidgetFactoryKwargs TypedDict "dynamically" based on the Pydantic model. This would at least get rid of the TypedDict (and the associated maintenance cost of keeping it in line with any changes to the fields of the Pydantic model). But building a TypedDict dynamically is something explicitly not supported for static type checking.

The Widget model, while technically dynamic, can be assumed to be static (no weird monkey-patching of my models after defining them).

Let's assume I have a pydantic model, such as this Widget:

models.py

from pydantic import BaseModel


class Widget(BaseModel):
    name: str
    value: float

When writing tests (using pytest), which use this Widget, I frequently want to be able to create widgets on the fly, with some default values for its fields (which I do not want to set as default values in general, ie on the model, because they are only meant as default values for tests), and potentially some fields being set to certain values.

For this I currently have this clunky construct in my conftest.py file:

conftest.py

from typing import NotRequired, Protocol, TypedDict, Unpack

import pytest

from models import Widget


class WidgetFactoryKwargs(TypedDict):
    name: NotRequired[str]
    value: NotRequired[float]


class WidgetFactory(Protocol):
    def __call__(self, **kwargs: Unpack[WidgetFactoryKwargs]) -> Widget: ...


@pytest.fixture
def widget_factory() -> WidgetFactory:
    def _widget_factory(**kwargs: Unpack[WidgetFactoryKwargs]) -> Widget:
        defaults = WidgetFactoryKwargs(name="foo", value=42)
        kwargs = defaults | kwargs
        return Widget(**kwargs)

    return _widget_factory

This gives the type checker the ability to check if I am using the factory correctly in my tests and gives my IDE autocompletion powers:

test_widgets.py

from typing import assert_type

from conftest import WidgetFactory


def test_widget_creation(widget_factory: WidgetFactory) -> None:
    widget = widget_factory()
    assert_type(widget, Widget)        # during type checking
    assert isinstance(widget, Widget)  # during run time
    assert widget.name == "foo"
    assert widget.value == 42

    widget = widget_factory(name="foobar")
    assert widget.name == "foobar"
    assert widget.value == 42

    widget = widget_factory(value=1337)
    assert widget.name == "foo"
    assert widget.value == 1337

    widget = widget_factory(name="foobar", value=1337)
    assert widget.name == "foobar"
    assert widget.value == 1337

    widget = widget_factory(mode="maintenance")  # type checker error

(the actual tests are of course more involved and use the widget in some other way)

Question:

Is there a better way to achieve this type safety? Ideally, I could build the WidgetFactoryKwargs TypedDict "dynamically" based on the Pydantic model. This would at least get rid of the TypedDict (and the associated maintenance cost of keeping it in line with any changes to the fields of the Pydantic model). But building a TypedDict dynamically is something explicitly not supported for static type checking.

The Widget model, while technically dynamic, can be assumed to be static (no weird monkey-patching of my models after defining them).

• python
• pytest
• python-typing
• pydantic

Share Improve this question Follow edited Nov 20, 2024 at 11:27 InSync 10.9k44 gold badges1717 silver badges5656 bronze badges asked Nov 20, 2024 at 10:17 GraipherGraipher 7,2062929 silver badges4848 bronze badges 1

• "Is there a better way to achieve this type safety?": Unfortunately, no. Transforming a TypedDict is quite a common problem, yet there exists no good solution for that so far. – InSync Commented Nov 20, 2024 at 11:56

Add a comment  | 

1 Answer 1

Sorted by: Reset to default 0

Basically, you're looking for a way of only needing to define the fields once, rather than repeating the same information in another TypedDict.

The main problem is to make an appropriate signature for WidgetFactory.__call__ such that it fulfils the following:

• Only accepts keyword arguments;
• Only accepts arguments from Widget.__init__;
• Doesn't need all arguments from Widget.__init__; omitting some or all of them is OK.

TypedDict, unfortunately, won't get you very far due to the lack of support for generics or key-value type manipulation. However, up-to-date versions of type-checkers should support @functools.partial to help us to do the same thing.

The exact behaviour may vary slightly between type-checkers, but mypy and pyright should accept some variation of the following strategy, the results of which is demonstrated here:

• mypy Playground
• pyright Playground

These two playground snippets differ slightly due to slightly different behaviours between the type-checkers.

---

First, some boilerplate to help us define WidgetFactory and pytest.fixture(widget_factory):

from __future__ import annotations

import typing_extensions as t

if t.TYPE_CHECKING:
    import collections.abc as cx

    R_co = t.TypeVar("R_co", covariant=True)
    PartialConstructorT = t.TypeVar("PartialConstructorT", bound=cx.Callable[..., t.Any])

    class KWOnlySignature(t.Protocol[R_co]):
        def __call__(self, /, **kwargs: t.Any) -> R_co: ...

def with_kw_only_partial_constructor_signature(
    f: PartialConstructorT, /
) -> cx.Callable[[cx.Callable[..., R_co]], PartialConstructorT | KWOnlySignature[R_co]]:
    return lambda _: f

Now, for the actual definitions. I'll use @dataclasses.dataclass as a substitute for pydantic.BaseModel:

import functools
from dataclasses import dataclass

@dataclass
class Widget:
    name: str
    value: float

class WidgetFactory(t.Protocol):
    @with_kw_only_partial_constructor_signature(functools.partial(Widget, name="foo", value=42))
    def __call__(self, /, **kwargs: t.Any) -> Widget: ...

    @classmethod
    def get_keyword_defaults(cls, /) -> dict[str, t.Any]: ...

@pytest.fixture
def widget_factory() -> WidgetFactory:
    def _widget_factory(**kwargs: t.Any) -> Widget:
        defaults = WidgetFactory.get_keyword_defaults()
        kwargs = defaults | kwargs
        return Widget(**kwargs)

    return _widget_factory  # type: ignore[return-value]

Here, we use functools.partial to simultaneously (1) define your expected default keyword arguments, and (2) manipulate the signature of WidgetFactory.__call__ so that the defaults do not need to be provided. If you try to pass incorrect arguments at this stage, it'll already show type-checker warnings:

class WidgetFactory(t.Protocol):
    # Type-checker error: no argument `Name`
    @with_kw_only_partial_constructor_signature(functools.partial(Widget, Name="foo"))
    def __call__(self, /, **kwargs: t.Any) -> Widget: ...

    @classmethod
    def get_keyword_defaults(cls, /) -> dict[str, t.Any]: ...

(Implementation of classmethod(get_keyword_defaults) is left as an exercise to the reader.)

Now, your tests involving WidgetFactory should now be invoked in a type-safe manner:

def test_widget_creation(widget_factory: WidgetFactory) -> None:
    widget_factory()                           # OK, no arguments required
    widget_factory(name="foobar")              # OK, one argument provided
    widget_factory(value=1337)                 # OK, one argument provided
    widget_factory(name="foobar", value=1337)  # OK, both arguments provided
    widget_factory("foobar", value=1337)       # Error: positional arguments not allowed
    widget_factory(mode="maintenance")         # Error: unrecognised argument