python_arango_ogm.db.pao_fields

  1from abc import ABC, abstractmethod
  2from enum import StrEnum, auto
  3from typing import Dict
  4
  5
  6class FieldTypeEnum(StrEnum):
  7    """
  8    Field Type Enum, used to specify field type in certain situations:
  9    """
 10
 11    ARRAY = auto()
 12    FLOAT = auto()
 13    INT = auto()
 14    STRING = auto()
 15
 16
 17class Field(ABC):
 18    def __init__(
 19            self,
 20            field_type: FieldTypeEnum,
 21            index_name: str = None,
 22            required: bool = False,
 23            unique: bool = False,
 24            minimum: float = None,
 25            maximum: float = None
 26    ):
 27        if unique and not index_name:
 28            raise ValueError(f"unique attribute requires an index_name for field {self.__class__.__name__}")
 29        self.field_type = field_type
 30        self.index_name = index_name
 31        self.unique = unique
 32        self.required = required
 33        self.minimum = minimum
 34        self.maximum = maximum
 35
 36    @abstractmethod
 37    def build_schema_properties(self) -> Dict[str, any]:
 38        pass
 39
 40
 41class FloatField(Field):
 42    def __init__(self,
 43                 index_name: str = None,
 44                 required: bool = False,
 45                 unique: bool = False,
 46                 minimum: float = None,
 47                 maximum: float = None):
 48        super().__init__(FieldTypeEnum.FLOAT, index_name=index_name, unique=unique, required=required, minimum=minimum,
 49                         maximum=maximum)
 50
 51    def build_schema_properties(self) -> Dict[str, any]:
 52        props = {'type': 'number'}
 53        if self.minimum is not None:
 54            props['minimum'] = self.minimum
 55
 56        if self.maximum is not None:
 57            props['maximum'] = self.maximum
 58        return props
 59
 60
 61class IntField(Field):
 62    def __init__(self,
 63                 index_name: str = None,
 64                 unique: bool = False,
 65                 required: bool = False,
 66                 minimum=None,
 67                 maximum=None):
 68        super().__init__(FieldTypeEnum.INT, index_name=index_name, unique=unique, required=required, minimum=minimum,
 69                         maximum=maximum)
 70
 71    def build_schema_properties(self) -> Dict[str, any]:
 72        props = {'type': 'number', 'multipleOf': 1}
 73        if self.minimum is not None:
 74            props['minimum'] = self.minimum
 75
 76        if self.maximum is not None:
 77            props['maximum'] = self.maximum
 78
 79        return props
 80
 81
 82class StrField(Field):
 83    def __init__(self,
 84                 index_name: str = None,
 85                 unique: bool = False,
 86                 required: bool = False,
 87                 minimum=None,
 88                 maximum=None):
 89        super().__init__(FieldTypeEnum.STRING, index_name=index_name, unique=unique, required=required, minimum=minimum,
 90                         maximum=maximum)
 91
 92    def build_schema_properties(self) -> Dict[str, any]:
 93        props = {'type': 'string'}
 94        if self.minimum is not None:
 95            props['minimum'] = self.minimum
 96
 97        if self.maximum is not None:
 98            props['maximum'] = self.maximum
 99
100        return props
101
102
103class ArrayField(Field):
104    def __init__(self, array_type: FieldTypeEnum, required: bool = False, minimum=None, maximum=None):
105        self.array_type = array_type
106        self.minimum = minimum
107        self.maximum = maximum
108        super().__init__(FieldTypeEnum.ARRAY, required=required)
109
110    def build_schema_properties(self) -> Dict[str, any]:
111        inner_field_type = 'number' if self.array_type in [FieldTypeEnum.INT, FieldTypeEnum.FLOAT] else 'string'
112
113        schema = {
114            'type': 'array',
115            'items': {
116                'type': inner_field_type
117            }
118        }
119
120        if self.minimum is not None:
121            schema['items']['minimum'] = self.minimum
122
123        if self.maximum is not None:
124            schema['items']['maximum'] = self.maximum
125
126        return schema