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

user_15024002 0
Viewed 1459

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?

python-dataclasses python-class python
4 Answers

This scenario is exactly what enums are made for.

>>> from enum import Enum
>>> class Material(Enum):
...     gold = 1
...     silver = 2
...     platinum = 3
...
>>> Material.diamond
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\...\enum.py", line 341, in __getattr__
    raise AttributeError(name) from None
AttributeError: diamond
>>> Material.gold
<Material.gold: 1>
user_14572690
  1. You should create ok_materials as a set instead.
  2. According to PEP-8, you should put the not before self.material
  3. ok_materials should rather be a class attribute (at least from what I see of your class)
  4. If you don't want the user to modify ok_materials, you can put an underscore at the beggining. In terms of dataclass specificities, I am clearly not familiar enough with this class to give an answer.
user_151217810

I might be a bit late to the party, but if you only want to validate for that simple scenario you can try chili library: https://github.com/kodemore/chili

Chili has a type integrity validation built-in which will work for your scenario, consider the following example:

from enum import Enum
from chili import init_dataclass

class Material(Enum):
    SILVER = "silver"
    GOLD = "gold"
    PLATINUM = "platinum"


@dataclass
class Jewel:
    kind: str
    material: Material

j1 = init_dataclass({"kind": "test", "material": "silver"}, Jewel)  # this works fine

j2 = init_dataclass({"kind": "test", "material": "test"}, Jewel)  # this will fail with a ValueError
user_28101300

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.

user_102375060
Related