André Ricardo
André Ricardo

Reputation: 3172

How to type hint builder pattern in Python?

Would like to add type hinting to def make(self): from in the class AggregateMaker so that the code in the tests test_fruit and test_tea would autocomplete the Fruit or Tea methods/properties rather than returning None

Is this possible in Python 3.10?

from dataclasses import dataclass

@dataclass
class Fruit:
    name: str
    smell: str

@dataclass
class Tea:
    name: str
    hot: bool

class AggregateMaker():
    _fields: dict

    @classmethod
    def new(cls, **fields):
        return cls(fields=None).with_(**fields)

    ###
    #  How to type hint in here to return Fruit or Tea?
    ###
    def make(self):
        return self._make(self._fields)

    def with_(self, **overrides):
        copy = dict(self._fields)
        for name, value in overrides.items():
            copy[name] = value
        return type(self)(copy)

class FruitMaker(AggregateMaker):
    def __init__(self, fields):
        if fields is None:
            fields = {
                "name": None,
                "smell": None,
            }
        self._fields = fields

    def _make(self, fields) -> Fruit:
        return Fruit(**fields)

class TeaMaker(AggregateMaker):
    def __init__(self, fields):
        if fields is None:
            fields = {
                "name": None,
                "hot": None,
            }
        self._fields = fields

    def _make(self, fields) -> Tea:
        return Tea(**fields)

def test_fruit():
    durian = FruitMaker.new().with_(name="Durian").with_(smell="Strong").make()
    assert durian.name == "Durian"
    assert durian.smell == "Strong"
    assert type(durian) is Fruit

def test_tea():
    camomile = TeaMaker.new(name="Camomile", hot=True).make()
    assert type(camomile) is Tea

Upvotes: 2

Views: 1099

Answers (1)

mwchase
mwchase

Reputation: 811

I typed as much of it as I felt was reasonable, but there are still gaps.

I feel like it'd usually make sense to replace most of this with, like, prototype objects, and calls to dataclasses.replace.

(From context elsewhere, I know that isn't practical in the near term.)

from dataclasses import dataclass
from typing import Any, Generic, Type, TypeVar

T = TypeVar("T")
TMaker = TypeVar("TMaker", bound="AggregateMaker[Any]")

@dataclass
class Fruit:
    name: str
    smell: str

@dataclass
class Tea:
    name: str
    hot: bool

class AggregateMaker(Generic[T]):
    _fields: dict[str, Any]
    
    def __init__(self, fields: dict[str, Any] | None) -> None:
        ...

    @classmethod
    def new(cls: Type[TMaker], **fields: Any) -> TMaker:
        return cls(fields=None).with_(**fields)

    def make(self) -> T:
        return self._make(self._fields)
        
    def _make(self, fields: dict[str, Any]) -> T:
        ...

    def with_(self: TMaker, **overrides: Any) -> TMaker:
        copy = dict(self._fields)
        for name, value in overrides.items():
            copy[name] = value
        return type(self)(copy)

class FruitMaker(AggregateMaker[Fruit]):
    def __init__(self, fields: dict[str, Any]):
        if fields is None:
            fields = {
                "name": None,
                "smell": None,
            }
        self._fields = fields

    def _make(self, fields: dict[str, Any]) -> Fruit:
        return Fruit(**fields)

class TeaMaker(AggregateMaker[Tea]):
    def __init__(self, fields: dict[str, Any]):
        if fields is None:
            fields = {
                "name": None,
                "hot": None,
            }
        self._fields = fields

    def _make(self, fields: dict[str, Any]) -> Tea:
        return Tea(**fields)

def test_fruit() -> None:
    durian = FruitMaker.new().with_(name="Durian").with_(smell="Strong").make()
    assert durian.name == "Durian"
    assert durian.smell == "Strong"
    assert type(durian) is Fruit

def test_tea() -> None:
    camomile = TeaMaker.new(name="Camomile", hot=True).make()
    assert type(camomile) is Tea

Upvotes: 3

Related Questions