From 8737897ce9dd0d198d2abcf2e7ab2640d390449e Mon Sep 17 00:00:00 2001
From: Peter Vacho <p_vacho@utb.cz>
Date: Wed, 25 Dec 2024 20:24:11 +0100
Subject: [PATCH] Add color validation for categories

---
 pyproject.toml        |  1 +
 requirements-dev.lock |  4 ++++
 requirements.lock     |  4 ++++
 src/api/categories.py | 43 +++++++++++++++++++++++++++----------------
 4 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 785458d..d001cec 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -15,6 +15,7 @@ dependencies = [
     "python-jose>=3.3.0",
     "python-multipart>=0.0.17",
     "bcrypt>=4.2.1",
+    "pydantic-extra-types>=2.10.1",
 ]
 readme = "README.md"
 requires-python = ">= 3.12"
diff --git a/requirements-dev.lock b/requirements-dev.lock
index 807ae28..96d6659 100644
--- a/requirements-dev.lock
+++ b/requirements-dev.lock
@@ -72,8 +72,11 @@ pydantic==2.10.2
     # via event-management
     # via fastapi
     # via lazy-model
+    # via pydantic-extra-types
 pydantic-core==2.27.1
     # via pydantic
+pydantic-extra-types==2.10.1
+    # via event-management
 pymongo==4.9.2
     # via motor
 python-decouple==3.8
@@ -101,6 +104,7 @@ typing-extensions==4.12.2
     # via fastapi
     # via pydantic
     # via pydantic-core
+    # via pydantic-extra-types
 uvicorn==0.32.1
     # via event-management
 virtualenv==20.28.0
diff --git a/requirements.lock b/requirements.lock
index 4e439d7..9ba7e25 100644
--- a/requirements.lock
+++ b/requirements.lock
@@ -56,8 +56,11 @@ pydantic==2.10.2
     # via event-management
     # via fastapi
     # via lazy-model
+    # via pydantic-extra-types
 pydantic-core==2.27.1
     # via pydantic
+pydantic-extra-types==2.10.1
+    # via event-management
 pymongo==4.9.2
     # via motor
 python-decouple==3.8
@@ -83,5 +86,6 @@ typing-extensions==4.12.2
     # via fastapi
     # via pydantic
     # via pydantic-core
+    # via pydantic-extra-types
 uvicorn==0.32.1
     # via event-management
diff --git a/src/api/categories.py b/src/api/categories.py
index 843bf64..35fd959 100644
--- a/src/api/categories.py
+++ b/src/api/categories.py
@@ -3,7 +3,8 @@ from typing import Annotated, final
 
 from beanie import PydanticObjectId
 from fastapi import APIRouter, Body, HTTPException, Response, status
-from pydantic import BaseModel
+from pydantic import BaseModel, StringConstraints, field_validator
+from pydantic_extra_types.color import Color
 
 from src.api.auth.dependencies import LoggedInDep
 from src.db.models.category import Category
@@ -21,13 +22,26 @@ categories_router = APIRouter(tags=["Categories"], prefix="/categories", depende
 base_router = APIRouter(tags=["Categories"], dependencies=[LoggedInDep])
 
 
+class _BaseCategoryData(BaseModel):
+    """Base class for all category data classes."""
+
+    name: Annotated[str, StringConstraints(max_length=50)]
+    color: Color
+
+    @field_validator("color", mode="after")
+    @classmethod
+    def validate_color(cls, value: Color) -> Color:
+        """Validate the color."""
+        if len(value.as_rgb_tuple()) == 4:
+            raise ValueError("Alpha channel is not allowed in colors")
+        return value
+
+
 @final
-class CategoryData(BaseModel):
+class CategoryData(_BaseCategoryData):
     """Data about a category sent to the user."""
 
     owner_user_id: PydanticObjectId
-    name: str
-    color: str
     created_at: datetime
 
     @classmethod
@@ -37,26 +51,23 @@ class CategoryData(BaseModel):
             raise ValueError("Got a category without id")
 
         return cls(
-            owner_user_id=category.id,
             name=category.name,
-            color=category.color,
+            color=Color(category.color),
+            owner_user_id=category.id,
             created_at=category.created_at,
         )
 
 
 @final
-class CategoryCreateData(BaseModel):
+class CategoryCreateData(_BaseCategoryData):
     """Data necessary to create a new category.
 
     This structure is intended to be used for POST & PUT requests.
     """
 
-    name: str
-    color: str
-
     async def create_category(self, user: User) -> Category:
         """Create a new category in the database."""
-        cat = Category(user=user, name=self.name, color=self.color)
+        cat = Category(user=user, name=self.name, color=self.color.as_hex(format="long"))
         return await cat.create()
 
     async def update_category(self, category: Category) -> Category:
@@ -67,7 +78,7 @@ class CategoryCreateData(BaseModel):
             category.name = self.name
             updated = True
         if category.color != self.color:
-            category.color = self.color
+            category.color = self.color.as_hex(format="long")
             updated = True
 
         if updated:
@@ -76,14 +87,14 @@ class CategoryCreateData(BaseModel):
 
 
 @final
-class PartialCategoryUpdateData(BaseModel):
+class PartialCategoryUpdateData(_BaseCategoryData):
     """Data necessary to perform a partial update of the category.
 
     This structure is intended to be used for PATCH requests.
     """
 
-    name: str | None = None
-    color: str | None = None
+    name: Annotated[str, StringConstraints(max_length=50)] | None = None  # pyright: ignore[reportIncompatibleVariableOverride]
+    color: Color | None = None  # pyright: ignore[reportIncompatibleVariableOverride]
 
     async def update_category(self, category: Category) -> Category:
         """Update an existing category, overwriting it with data that were specified."""
@@ -93,7 +104,7 @@ class PartialCategoryUpdateData(BaseModel):
             category.name = self.name
             updated = True
         if self.color and category.color != self.color:
-            category.color = self.color
+            category.color = self.color.as_hex(format="long")
             updated = True
 
         if updated: