orm/base/enum.py

0001 #!/usr/bin/env python
0002 #
0003 # Licensed under the Apache License, Version 2.0 (the "License");
0004 # You may not use this file except in compliance with the License.
0005 # You may obtain a copy of the License at
0006 # http://www.apache.org/licenses/LICENSE-2.0OA
0007 #
0008 # Authors:
0009 # - Wen Guan, <wen.guan@cern.ch>, 2019
0010
0011
0012 """
0013 handling sqlachelmy enum types.
0014
0015 Borrowed from:
0016 https://github.com/rucio/rucio/blob/master/lib/rucio/db/sqla/enum.py
0017 """
0018
0019 import uuid
0020
0021 from sqlalchemy.types import SchemaType, TypeDecorator, Enum
0022
0023 from idds.common.exceptions import InvalidDatabaseType
0024
0025
0026 class EnumSymbol(object):
0027     """Define a fixed symbol tied to a parent class."""
0028
0029     def __init__(self, cls_, name, value, description):
0030         self.cls_ = cls_
0031         self.name = name
0032         self.value = value
0033         self.description = description
0034
0035     def __reduce__(self):
0036         """Allow unpickling to return the symbol
0037         linked to the DeclEnum class."""
0038         return getattr, (self.cls_, self.name)
0039
0040     def __iter__(self):
0041         return iter([self.value, self.description])
0042
0043     def __repr__(self):
0044         return "%s" % self.description
0045
0046
0047 class EnumMeta(type):
0048     """Generate new DeclEnum classes."""
0049
0050     def __init__(cls, classname, bases, dict_):  # pylint: disable=E0101
0051         cls._reg = reg = cls._reg.copy()
0052         cls._syms = syms = cls._syms.copy()
0053         for k, v in dict_.items():
0054             if isinstance(v, tuple):
0055                 sym = reg[v[0]] = syms[v[1]] = EnumSymbol(cls, k, *v)
0056                 setattr(cls, k, sym)
0057         return type.__init__(cls, classname, bases, dict_)
0058
0059     def __iter__(cls):
0060         return iter(cls._reg.values())
0061
0062
0063 class DeclEnum(object):
0064     """Declarative enumeration."""
0065
0066     __metaclass__ = EnumMeta
0067     _reg = {}
0068     _syms = {}
0069
0070     @classmethod
0071     def from_string(cls, value):
0072         try:
0073             return cls._reg[value]
0074         except KeyError:
0075             raise ValueError("Invalid value for %r: %r" % (cls.__name__, value))
0076
0077     @classmethod
0078     def from_sym(cls, value):
0079         try:
0080             return cls._syms[value.upper()]
0081         except KeyError:
0082             raise ValueError("Invalid value for %r: %r" % (cls.__name__, value))
0083
0084     @classmethod
0085     def values(cls):
0086         return cls._reg.keys()
0087
0088     @classmethod
0089     def db_type(cls, name=None, default=None):
0090         return DeclEnumType(enum=cls, name=name, default=default)
0091
0092
0093 class DeclEnumType(SchemaType, TypeDecorator):
0094
0095     def __init__(self, enum, name=None, default=None):
0096         self.enum = enum
0097         if name is None:
0098             self.impl = Enum(*enum.values(), name='RUCIO_ENUM_' + str(uuid.uuid4())[:6])
0099         else:
0100             self.impl = Enum(*enum.values(), name=name)
0101
0102     def _set_parent_with_dispatch(self, parent):
0103         TypeDecorator._set_parent_with_dispatch(self, parent)
0104         SchemaType._set_parent_with_dispatch(self, parent)
0105
0106     def copy(self):
0107         return DeclEnumType(self.enum)
0108
0109     def process_bind_param(self, value, dialect):
0110         try:
0111             if value is None:
0112                 return None
0113             return value.value
0114         except AttributeError:
0115             raise InvalidDatabaseType('Invalid value/type %s for %s' % (value, self.enum))
0116
0117     def process_result_value(self, value, dialect):
0118         if value is None:
0119             return None
0120         return self.enum.from_string(value.strip())