Reputation: 3172
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
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