A more pythonic way to define an enum with dynamic members

前端 未结 3 2206
爱一瞬间的悲伤
爱一瞬间的悲伤 2021-01-06 13:32

I needed to create an enum to represent the ISO country codes. The country code data comes from a json file which can be obtained from: https://github.com/lukes/ISO-3166-Cou

相关标签:
3条回答
  • 2021-01-06 14:05

    Yes, there is a way to define the enum using the alternate declaration syntax you want. It works by hiding the "ugly" code in a metaclass derived from enum.EnumMeta. If you wished, it would also be possible to define the choices() class method there, too.

    import enum
    import json
    
    class CountryCodeMeta(enum.EnumMeta):
        def __new__(metacls, cls, bases, classdict):
            data = classdict['data']
            names = [(country['alpha-2'], int(country['country-code'])) for country in data]
    
            temp = type(classdict)()
            for name, value in names:
                temp[name] = value
    
            excluded = set(temp) | set(('data',))
            temp.update(item for item in classdict.items() if item[0] not in excluded)
    
            return super(CountryCodeMeta, metacls).__new__(metacls, cls, bases, temp)
    
    class CountryCode(enum.Enum, metaclass=CountryCodeMeta):
        data = json.load(open('slim-2.json'))
    
        @classmethod
        def choices(cls):
            return ((member.value, name) for name, member in cls.__members__.items())
    
    0 讨论(0)
  • 2021-01-06 14:11

    How about this?

    data = json.load(open('slim-2.json'))
    CountryCode = enum.Enum('CountryCode', [
        (x['alpha-2'], int(x['country-code'])) for x in data
    ])
    CountryCode._names = {x['alpha-2']: x['name'] for x in data}
    CountryCode.__str__ = lambda self: self._names[self.name]
    CountryCode.choices = lambda: ((e.value, e.name) for e in CountryCode)
    
    • Replaced [...data[i]... for i in range(len(data))] with [...x... for x in data]; You can itearte sequence (list, data in the code) without using indexes.
    • Used CountryCode.attr = ... consistently; instead of mixing CountryCode.attr = ... and setattr(CountryCode, 'attr', ...).
    0 讨论(0)
  • 2021-01-06 14:18

    Update

    Using JSONEnum at the bottom of When should I subclass EnumMeta instead of Enum?, you can do this:

    class Country(JSONEnum):
        _init_ = 'abbr code country_name'  # remove if not using aenum
        _file = 'some_file.json'
        _name = 'alpha-2'
        _value = {
                1: ('alpha-2', None),
                2: ('country-code', lambda c: int(c)),
                3: ('name', None),
                }
    

    Original Answer

    It looks like you are trying to keep track of three pieces of data:

    • country name
    • country code
    • country 2-letter abbreviaton

    You should consider using a technique inspired by a namedtuple mixin as illustrated in this answer:


    The stdlib way

    We'll need a base class to hold the behavior:

    from enum import Enum
    import json
    
    class BaseCountry(Enum):
    
        def __new__(cls, record):
            member = object.__new__(cls)
            member.country_name = record['name']
            member.code = int(record['country-code'])
            member.abbr = record['alpha-2']
            member._value_ = member.abbr, member.code, member.country_name
            if not hasattr(cls, '_choices'):
                cls._choices = {}
            cls._choices[member.code] = member.country_name
            cls._choices[member.abbr] = member.country_name
            return member                
    
        def __str__(self):
            return self.country_name
    
        @classmethod
        def choices(cls):
            return cls._choices.copy()
    

    Then we can use that to create the actual Country class:

    Country = BaseCountry(
            'Country',
            [(rec['alpha-2'], rec) for rec in json.load(open('slim-2.json'))],
            )
    

    The aenum way 1 2

    from aenum import Enum, MultiValue
    import json
    
    class Country(Enum, init='abbr code country_name', settings=MultiValue):
    
        _ignore_ = 'this country'  # do not add these names as members
    
        # create members
        this = vars()
        for country in json.load(open('slim-2.json')):
            this[country['alpha-2']] = (
                    country['alpha-2'],
                    int(country['country-code']),
                    country['name'],
                    )
    
        # return a dict of choices by abbr or country code to name
        @classmethod
        def choices(cls):
            mapping = {}
            for member in cls:
                mapping[member.code] = member.name
                mapping[member.abbr] = member.name
            return mapping
    
        # have str() print just the country name
        def __str__(self):
            return self.country_name
    

    While I included the choices method, you may not need it:

    >>> Country('AF')
    <Country.AF: ('AF', 4, 'Afghanistan')>
    
    >>> Country(4)
    <Country.AF: ('AF', 4, 'Afghanistan')>
    
    >>> Country('Afghanistan')
    <Country.AF: ('AF', 4, 'Afghanistan')>
    

    1 Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.

    2 This requires aenum 2.0.5+.

    0 讨论(0)
提交回复
热议问题