rucio.db.enum — Rucio 0.3.8 documentation

# Source code for rucio.db.enum

    # Copyright European Organization for Nuclear Research (CERN)
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # You may not use this file except in compliance with the License.
    # You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
    #
    # Authors:
    # - Vincent Garonne, <vincent.garonne@cern.ch>, 2013
    # - Mario Lassnig, <mario.lassnig@cern.ch>, 2013-2014

    '''
    Class to handle enum type with sqlachelmy.
    ref. http://techspot.zzzeek.org/2011/01/14/the-enum-recipe/

    '''

    import uuid

    from sqlalchemy.types import SchemaType, TypeDecorator, Enum

    from rucio.common.exception import InvalidType

    class EnumSymbol(object):
        """Define a fixed symbol tied to a parent class."""

        def __init__(self, cls_, name, value, description):
            self.cls_ = cls_
            self.name = name
            self.value = value
            self.description = description

        def __reduce__(self):
            """Allow unpickling to return the symbol
            linked to the DeclEnum class."""
            return getattr, (self.cls_, self.name)

        def __iter__(self):
            return iter([self.value, self.description])

        def __repr__(self):
            return "%s" % self.description

    class EnumMeta(type):
        """Generate new DeclEnum classes."""

        def __init__(cls, classname, bases, dict_):
            cls._reg = reg = cls._reg.copy()
            cls._syms = syms = cls._syms.copy()
            for k, v in dict_.items():
                if isinstance(v, tuple):
                    sym = reg[v[0]] = syms[v[1]] = EnumSymbol(cls, k, *v)
                    setattr(cls, k, sym)
            return type.__init__(cls, classname, bases, dict_)

        def __iter__(cls):
            return iter(cls._reg.values())

    class DeclEnum(object):
        """Declarative enumeration."""

        __metaclass__ = EnumMeta
        _reg = {}
        _syms = {}

        @classmethod
        def from_string(cls, value):
            try:
                return cls._reg[value]
            except KeyError:
                raise ValueError("Invalid value for %r: %r" % (cls.__name__, value))

        @classmethod
        def from_sym(cls, value):
            try:
                return cls._syms[value.upper()]
            except KeyError:
                raise ValueError("Invalid value for %r: %r" % (cls.__name__, value))

        @classmethod
        def values(cls):
            return cls._reg.keys()

        @classmethod
        def db_type(cls, name=None, default=None):
            return DeclEnumType(enum=cls, name=name, default=default)

    class DeclEnumType(SchemaType, TypeDecorator):

        def __init__(self, enum, name=None, default=None):
            self.enum = enum

            # Workaround SQLAlchemy not propagating Enum names properly

            if name is None:
                self.impl = Enum(*enum.values(), name='RUCIO_ENUM_' + str(uuid.uuid4())[:6])
            else:
                self.impl = Enum(*enum.values(), name=name)

        def _set_table(self, table, column):
            self.impl._set_table(table, column)

        def copy(self):
            return DeclEnumType(self.enum)

        def process_bind_param(self, value, dialect):
            try:
                if value is None:
                    return None
                return value.value
            except AttributeError:
                raise InvalidType('Invalid value/type %s for %s' % (value, self.enum))

        def process_result_value(self, value, dialect):
            if value is None:
                return None
            return self.enum.from_string(value.strip())