[Numpy-discussion] Enum type

Dag Sverre Seljebotn d.s.seljebotn at astro.uio.no
Tue Jan 3 14:29:44 EST 2012

On 01/03/2012 06:46 PM, Ognen Duzlevski wrote:
> Hello,
> I am playing with adding an enum dtype to numpy (to get my feet wet in
> numpy really). I have looked at the
> https://github.com/martinling/numpy_quaternion and I feel comfortable
> with my understanding of adding a simple type to numpy in technical
> terms.
> I am mostly a C programmer and have programmed in Python but not at
> the level where my code wcould be considered "pretty" or maybe even
> "pythonic". I know enums from C and have browsed around a few python
> enum implementations online. Most of them use hash tables or lists to
> associate names to numbers - these approaches just feel "heavy" to me.

If you want the enum values to be stored efficiently (using 1, 2 or 
4-byte integers), and want a mapping between string names and such 
integers, then you need to map between them somehow, right? I.e., when 
printing the repr() of each element, you at least need a list in order 
to go from enum values to names (and that doesn't feel 'heavy' to me -- 
it's the minimal possible solution for the job!)

It's unclear whether you mean heavy on the CPU, in the API, in the C 
code, or whatever, so difficult to give more feedback.

As far as the API goes, you could probably do something like:

colors = np.enum(['red', 'green', 'blue'])
arr = np.asarray([colors.red, colors.red, colors.red, colors.blue])
assert arr[0] == colors.red
assert np.all(arr.view(np.int8) == [0, 0, 0, 2])

So the strings are only needed in the API in the constructor of the enum 
type. They are needed there though.

Dag Sverre

More information about the NumPy-Discussion mailing list