Search code examples
pythonpython-3.xormenumsponyorm

How can I store a Python Enum using Pony ORM?


Say I've got this simple little Pony ORM mapping here. The built-in Enum class is new as of Python 3.4, and backported to 2.7.

from enum import Enum

from pony.orm import Database, Required


class State(Enum):
    ready = 0
    running = 1
    errored = 2

if __name__ == '__main__':
    db = Database('sqlite', ':memory:', create_db=True)

    class StateTable(db.Entity):
        state = Required(State)

    db.generate_mapping(create_tables=True)

When I run the program, an error is thrown.

TypeError: No database converter found for type <enum 'State'>

This happens because Pony doesn't support mapping the enum type. Of course, the workaround here is to just store the Enum value, and provide a getter in Class StateTable to convert the value to the Enum once again. But this is tedious and error prone. I can also just use another ORM. Maybe I will if this issue becomes too much of a headache. But I would rather stick with Pony if I can.

I would much rather create a database converter to store the enum, like the error message is hinting at. Does anyone know how to do this?

UPDATE: Thanks to Ethan's help, I have come up with the following solution.

from enum import Enum

from pony.orm import Database, Required, db_session
from pony.orm.dbapiprovider import StrConverter


class State(Enum):
    ready = 0
    running = 1
    errored = 2

class EnumConverter(StrConverter):

    def validate(self, val):
        if not isinstance(val, Enum):
            raise ValueError('Must be an Enum.  Got {}'.format(type(val)))
        return val

    def py2sql(self, val):
        return val.name

    def sql2py(self, value):
        # Any enum type can be used, so py_type ensures the correct one is used to create the enum instance
        return self.py_type[value]

if __name__ == '__main__':
    db = Database('sqlite', ':memory:', create_db=True)

    # Register the type converter with the database
    db.provider.converter_classes.append((Enum, EnumConverter))

    class StateTable(db.Entity):
        state = Required(State)

    db.generate_mapping(create_tables=True)

    with db_session:
        s = StateTable(state=State.ready)
        print('Got {} from db'.format(s.state))

Solution

  • Excerpt from some random mailing list:

    2.2. CONVERTER METHODS

    Each converter class should define the following methods:

    class MySpecificConverter(Converter):
    
        def init(self, kwargs):
            # Override this method to process additional positional
            # and keyword arguments of the attribute
    
           if self.attr is not None:
                # self.attr.args can be analyzed here
                self.args = self.attr.args
    
            self.my_optional_argument = kwargs.pop("kwarg_name")
            # You should take all valid options from this kwargs
            # What is left in is regarded as unrecognized option
    
        def validate(self, val):
            # convert value to the necessary type (e.g. from string)
            # validate all necessary constraints (e.g. min/max bounds)
            return val
    
        def py2sql(self, val):
            # prepare the value (if necessary) to storing in the database
            return val
    
        def sql2py(self, value):
            # convert value (if necessary) after the reading from the db
            return val
    
        def sql_type(self):
            # generate corresponding SQL type, based on attribute options
            return "SOME_SQL_TYPE_DEFINITION"
    

    You can study the code of the existing converters to see how these methods are implemented.