Add Length class

pyplanner/models/items.py

@@ -3,6 +3,42 @@ from pyplanner.models import Base
 from sqlalchemy import Column, Integer, String, Text, ForeignKey
 from sqlalchemy.orm import relationship, backref
 from pydantic import BaseModel
+from enum import Enum
+from functools import total_ordering
+
+
+@total_ordering
+class Length(Enum):
+    MINUTES = 'mi'
+    SEVERAL_MINUTES = 'mi+'
+    HOURS = 'h'
+    SEVERAL_HOURS = 'h+'
+    DAYS = 'd'
+    SEVERAL_DAYS = 'd+'
+    WEEKS = 'w'
+    SEVERAL_WEEKS = 'w+'
+    MONTHS = 'mo'
+    SEVERAL_MONTHS = 'mo+'
+
+    def __add__(self, other):
+        all_values = list(Length.__members__.values())
+        if self == other and '+' not in self.value:
+            return Length(f'{self.value}+')
+        elif self.value in other.value:
+            index = all_values.index(other)
+            return all_values[index+1]
+        elif other.value in self.value:
+            index = all_values.index(self)
+            return all_values[index+1]
+        elif self > other:
+            return self
+        elif other > self:
+            return other
+        raise NotImplementedError()
+
+    def __lt__(self, other):
+        all_values = list(Length.__members__.values())
+        return all_values.index(self) < all_values.index(other)
 
 
 class SchemaItem(BaseModel):
@@ -15,7 +51,7 @@ class SchemaItem(BaseModel):
     sprint_id: int
 
     priority: str
-    length: str
+    length: Length = Length.DAYS
 
 
 class Item(SQLABase, Base):