5

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-Countries-with-Regional-Codes

So what I've done is:

data = json.load(open('slim-2.json'))
codes_list = [(data[i]['alpha-2'], int(data[i]['country-code']))
              for i in range(len(data))]

CountryCode = enum.Enum('CountryCode', codes_list,)

names_dict = {int(data[i]['country-code']):data[i]['name'] 
              for i in range(len(data))}
setattr(CountryCode, '_names', names_dict)

CountryCode.choices = classmethod(lambda cls:((member.value, name) 
                                  for name, member in cls.__members__.items()))
setattr(CountryCode, '__str__' ,lambda self: self.__class__._names[self.value])

This code snippet is frankly ugly. I looked at alternative ways to define the enum class but couldn't piece together a solution. Is there a way to define the enum in the following form:

class CountryCode(enum.Enum):

    data = json.load(open('slim-2.json'))
    # Some code to define the enum members

    @classmethod
    def choices(cls):
    # etc...

Any suggestions on how to do this?

Peter Pudaite
  • 406
  • 8
  • 18
  • The only way to programatically assign the members of an Enum is by using the `Enum()` initialiser, so unfortunately, there isn't a way to do this via the class declaration. And you can't subclass an Enum that already has members assigned. You can specify a mixin (the `type` parameter) to add behaviour, but that doesn't help much in your case. – dirkgroten Mar 29 '17 at 16:35
  • I suspected that... thanks. – Peter Pudaite Mar 29 '17 at 16:58
  • @dirkgroten: While this isn't possible with the stdlib `Enum`, it is possible with [`aenum`](https://pypi.python.org/pypi/aenum). – Ethan Furman Mar 29 '17 at 18:40
  • Learning a lot from this post :-) Both `aenum` and `metaclass` show that I was wrong... – dirkgroten Mar 29 '17 at 19:57
  • 1
    @PeterPudaite: If creating `Enum`s from json is common for you, check out [this answer](http://stackoverflow.com/q/43730305/208880) which details a `JSONEnum`. – Ethan Furman May 02 '17 at 04:53

3 Answers3

7

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+.

Ethan Furman
  • 63,992
  • 20
  • 159
  • 237
  • I wasn't aware of the aenum library. That does make it much clearer to understand what's going on. Just one thing. What is the purpose of `_ignore_ = 'this country'`? – Peter Pudaite Mar 29 '17 at 19:33
  • @PeterPudaite: Just added a comment to the code. `_ignore_` is a list of names that should not become `Enum` members; in this case, the `this` and `country` temporary variables. – Ethan Furman Mar 29 '17 at 19:54
  • OK gotcha. I'm looking at the aenum repo home page. Is there more documentation for aenum elsewhere? And thanks for creating the library! – Peter Pudaite Mar 29 '17 at 20:02
  • There is a `/doc` folder in the repo (should also be in the sdist) that has a doc with a bunch of the enhancements listed. – Ethan Furman Mar 29 '17 at 20:05
  • Ethan: In the interests of transparency and full disclosure, I think you should mention that **you're** the primary author of the stdlid `enum` module—I've already seen you not do it at least a couple of times today. – martineau Mar 29 '17 at 21:26
  • 1
    @martineau: [It's an interesting conundrum.](https://meta.stackoverflow.com/a/317309/208880) On the one hand, I want folks to trust that `aenum` is a solid package; on the other hand, I like getting kudos for a good library; on the gripping hand, I want up-votes because my answer is good, not because I'm the Python core developer that happened to write `Enum`, `enum34`, and `aenum`. By way of contrast, if you look at my [`dbf` answers](https://stackoverflow.com/search?q=%5Bdbf%5Duser%3A208880) you'll see that anytime I reference my `dbf` library I disclose that I'm the author. – Ethan Furman Mar 29 '17 at 21:55
  • It could be that because of your inside knowledge due to the fact you wrote _both_ modules, you have what some might consider an unfair advantage (especially with respect to the `aenum` module) answering questions related to them. BTW, while I have your attention: I think you should modify `enum._EnumDict` so its constructor supports arguments the same as normal dictionaries do. Doing so would be backwards-compatible with existing code. – martineau Mar 29 '17 at 22:10
  • @Ethan: I have no problem with you answering `enum`-related questions, just when you use it primarily as a means to promote another module you wrote. Making `_EnumDict` more "normal" would have made doing what I did via a metaclass easier—and after looking a the source code, I don't see any problem with making it so. I got the impression that exposing the `enum` metaclass was intentional so as to allow alternative declaration syntax from [an article](http://python-notes.curiousefficiency.org/en/latest/python3/enum_creation.html) I read recently. – martineau Mar 29 '17 at 22:23
  • 1
    @martineau: I'll look into those `_EnumDIct` changes, and also about making `_EnumDict` public, since it's difficult to work with `EnumMeta` without it. As far as extending `EnumMeta` -- sure, get the word out that it can be done -- but preferably as a secondary portion of the answer with the primary portion being the boring but safe `__new__` or plain old ordinary methods. PEP 435 `Enum`s are not ordinary classes, and `EnumMeta` is not a typical metaclass. As you've seen, it is easy to break them if unaware of the underlying complexities. – Ethan Furman Mar 30 '17 at 07:27
  • @Ethan: I'll look into it—will need to jog my memory regarding the details of our month-old conversation. – martineau Apr 30 '17 at 08:52
  • @martineau: Okay, having looked through the code and the reasons behind having `_EnumDict` private, I'm going to leave it that way. `_EnumDict` is not meant to stand alone, and any extra initialization is/will be done in the `__prepare__` method of `EnumMeta`. If you need another instance of it call `EnumMeta.__prepare__(...)` to get one. – Ethan Furman May 02 '17 at 05:04
  • @EthanFurman: I wasn't asking for `enum._EnumDict` to be made public, just for it behave more like the `dict` constructor, argument-wise. The main motivation was, as you can see in the code of [my answer](http://stackoverflow.com/a/43100224/355230) to this question, the derived `Enum` metaclass instantiates and populates a new one. It avoids breaking encapsulation by indirectly referring to the private type as `type(classdict)`. However, due to the differing constructor signatures, it can't be initialized as easily as a `dict` can be—thereby forcing that to be done explicitly and manually. – martineau May 02 '17 at 19:15
  • @EthanFurman: P.S. I'm not very familiar with Python 3's metaclass `__prepare__` method, but will now take a closer look and see whether/how it presents a better way to do things. – martineau May 02 '17 at 19:22
1

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', ...).
falsetru
  • 357,413
  • 63
  • 732
  • 636
1

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())
martineau
  • 119,623
  • 25
  • 170
  • 301
  • OK I've never used meta classes so this approach is new to me. I can understand to some degree what's going on in the first few lines but the block from `temp = type(....` until `temp.update(....` doesn't make sense to me. What is that block doing? – Peter Pudaite Mar 29 '17 at 19:45
  • 1
    @Peter: That section first creates an empty temporary instance of a dictionary-like `enum._EnumDict'` class to use in the `super()` call at the end. It first populates this with the names and values extracted from the JSON data. Then it creates a `set` of the things it doesn't want copied from the `classdict` argument it was passed, which includes any of the JSON defined names it just put in, plus the name `'data'`, which it knows was defined in the `CountryCode` class definition and shouldn't be an `enum` member. Lastly it copies anything else from `classdict` that hasn't been excluded. – martineau Mar 29 '17 at 20:24
  • ...continued: A nice thing about metaclasses is that they can do almost anything if you take the time to understand the nitty-gritty of how they work. For the `enum` class it's a little more involved because it already has a custom metaclass that needs to also be understood—but most of the time for "normal" classes, that won't be necessary. – martineau Mar 29 '17 at 20:31
  • Thanks. I've started to read up on it. I think I'll be making use of it for other things soon. – Peter Pudaite Mar 29 '17 at 20:40
  • Good to hear—but just remember, "with great power comes great responsibility" (Spider-Man). – martineau Mar 29 '17 at 21:17
  • @Ethan: Are you saying you think using a metaclass for this is worse than using the non-stdlib `aenum` module or are you just trying to promote it? – martineau Mar 29 '17 at 22:14
  • @Ethan For the case of enums I will probably go with your library solution. However I have other classes which need to be defined through configuration which would probably benefit from using a metaclass approach. I like to keep things simple but balance that with the need to be resilient to change or variation. – Peter Pudaite Mar 29 '17 at 22:34
  • 1
    @martineau: I created a question/answer about [deriving from `EnumMeta`](http://stackoverflow.com/q/43730305/208880). – Ethan Furman May 02 '17 at 04:57