ExternalCompilerError
ExternalCompilerError

Reputation: 109

Best way to implement a dict with a fixed key type (e.g. an Enum type)?

The Setting

I have a data object that comprises a number of properties, let's say:

FaceProperties( eye_color:EyeColorEnum, has_glasses:bool, 
                nose_length_mm:float, hair_style:str )

What happened so far

In principle, I would very much like to collect them in a dict, so that I can do all the nice things one can do with a dict (extend with new keys, iterate over, access by key...). However, an ordinary dict (and linter) would not notice when an unsupported key is used, e.g. due to some typo, especially with keys that are strings.

Due to these concerns, I started the project using an ordinary class with ordinary properties and implemented all the generic stuff (iteration etc) manually.

Are there better options?

I am currently refactoring the software, and started wondering if a map with a fixed key type (for which I could use a SupportedProperties Enum) would solve my problem better, and if there is such a thing in python.

Now, I can simply use the Enum values as keys with an ordinary dict and probably be reasonably fine, but it would be even nicer if there was some specialized dict class that would not even accept keys of other types.

Is there maybe some collection (preferably a builtin one) I could use for that? Or would I have to derive my own class from UserDict and do the type checking myself?

Upvotes: 2

Views: 1371

Answers (2)

Mahi
Mahi

Reputation: 21951

As others have mentioned in the comments, all you need is one check to ensure that the key is in a certain set of values:

class FixedKeyDict(dict):
    def __init__(self, *args, allowed_keys=(), **kwargs):
        super().__init__(*args, **kwargs)
        self.allowed_keys = allowed_keys

    def __setitem__(self, key, value):
        if key not in self.allowed_keys:
            raise KeyError('Key {} is not allowed'.format(key))
        super().__setitem__(key, value)

Now give it a list of the allowed values:

hair_dict = FixedKeyDict(allowed_keys=('bald', 'wavy', 'straight', 'mohawk'))

This only overrides the [] operator though, you want to subclass UserDict like you've said and override all the other methods too, but you get the idea. Regardless, you must implement it yourself, such thing doesn't exist in the standard library.

Upvotes: 2

AKX
AKX

Reputation: 169407

How about dataclasses (Python 3.7+)?

from enum import Enum
from dataclasses import dataclass, asdict


class EyeColorEnum(Enum):
    BLUE = 1
    RED = 2


@dataclass()
class FaceProperties:
    eye_color: EyeColorEnum = None
    has_glasses: bool = None
    nose_length_mm: float = None
    hair_style: str = None


fp = FaceProperties(eye_color=EyeColorEnum.BLUE)
print(fp)
print(asdict(fp))

outputs

FaceProperties(eye_color=<EyeColorEnum.BLUE: 1>, has_glasses=None, nose_length_mm=None, hair_style=None)
{'eye_color': <EyeColorEnum.BLUE: 1>, 'has_glasses': None, 'nose_length_mm': None, 'hair_style': None}

Upvotes: 2

Related Questions