Python dataclass: Check that a field is in list of possible values

I want a field in a Python dataclass to only have certain possible values. For instance, I want a Jewel dataclass, whose field material can only accept values "gold", "silver", and "platinum". I came up with this solution:

@dataclass
class Jewel:
    kind: str
    material: str
    
    def __post_init__(self):
        ok_materials = ["gold", "silver", "platinum"]
        if self.material not in ok_materials:
            raise ValueError(f"Only available materials: {ok_materials}.")

So, Jewel("earrings", "silver") would be ok, whereas Jewel("earrings", "plastic") would raise a ValueError.

Is there a better or more pythonic way to achieve this? Maybe something that relies on features from dataclasses module?

Solution

I would also propose the dataclass-wizard library, which enables type validation when loading JSON or dict data to a dataclass type, such as with the fromdict helper function for example.

To make it simple, you could still pass in a string but use typing.Literal to enforce strict value checks. Alternatively, you could define an Enum class to hold the possible values, in case that works a little better for you.

Here is an example using typing.Literal to restrict the possible values. In Python 3.7 or earlier, I would import Literal from the typing_extensions module instead.

from dataclasses import dataclass
from typing import Literal

from dataclass_wizard import fromdict
from dataclass_wizard.errors import ParseError


@dataclass
class Jewel:
    kind: str
    # ok_materials = ["gold", "silver", "platinum"]
    material: Literal['gold', 'silver', 'platinum']


print(fromdict(Jewel, {'kind': 'test', 'material': 'gold'}))
print()

# following should raise a `ParseError`, as 'copper' is not in
# the valid Literal values.
try:
    _ = fromdict(Jewel, {'kind': 'test', 'material': 'copper'})
except ParseError as e:
    print(e)

The output is nice and clean, with a clear error message displayed about what went wrong:

Jewel(kind='test', material='gold')

Failure parsing field `material` in class `Jewel`. Expected a type Literal, got str.
  value: 'copper'
  error: Value not in expected Literal values
  allowed_values: ['gold', 'silver', 'platinum']
  json_object: '{"kind": "test", "material": "copper"}'

Disclaimer: I am the creator and maintenor of this library.