From c31c3dea42fd400f6e769937125fbceae81210fe Mon Sep 17 00:00:00 2001
From: Donkie <daniel.cf.hultgren@gmail.com>
Date: Sun, 2 Apr 2023 15:48:30 +0200
Subject: [PATCH] Added API for using spool weight/length

---
 pyproject.toml             |  2 +-
 spoolson/api/v1/spool.py   | 37 +++++++++++++++++++++++++++
 spoolson/database/spool.py | 52 ++++++++++++++++++++++++++++++++++++--
 spoolson/math.py           | 19 ++++++++++++++
 4 files changed, 107 insertions(+), 3 deletions(-)
 create mode 100644 spoolson/math.py

diff --git a/pyproject.toml b/pyproject.toml
index c3be0aa..1317349 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,7 +8,7 @@ license = {text = "AGPL-3.0-only"}
 
 [tool.ruff]
 select = ["ALL"]
-ignore = ["A003", "D101", "D104", "S104", "TRY201", "TRY003", "EM101", "EM102"]
+ignore = ["A003", "D101", "D104", "D407", "S104", "TRY201", "TRY003", "EM101", "EM102"]
 line-length = 120
 target-version = "py39"
 
diff --git a/spoolson/api/v1/spool.py b/spoolson/api/v1/spool.py
index 3177b9e..39c8e7a 100644
--- a/spoolson/api/v1/spool.py
+++ b/spoolson/api/v1/spool.py
@@ -43,6 +43,11 @@ class SpoolUpdateParameters(SpoolParameters):
     filament_id: Optional[int] = Field(description="The ID of the filament type of this spool.")
 
 
+class SpoolUseParameters(BaseModel):
+    use_length: Optional[float] = Field(description="Length of filament to reduce by, in mm.")
+    use_weight: Optional[float] = Field(description="Filament weight to reduce by, in g.")
+
+
 @router.get("/")
 async def find(_filament: Union[int, None] = None) -> list[Spool]:
     return []
@@ -116,3 +121,35 @@ async def delete(
 ) -> Message:
     await spool.delete(db, spool_id)
     return Message(message="Success!")
+
+
+@router.put(
+    "/{spool_id}/use",
+    response_model=Spool,
+    responses={
+        400: {"model": Message},
+    },
+)
+async def use(  # noqa: ANN201
+    db: Annotated[AsyncSession, Depends(get_db_session)],
+    spool_id: int,
+    body: SpoolUseParameters,
+):
+    if body.use_weight is not None and body.use_length is not None:
+        return JSONResponse(
+            status_code=400,
+            content={"message": "Only specify either use_weight or use_length."},
+        )
+
+    if body.use_weight is not None:
+        db_item = await spool.use_weight(db, spool_id, body.use_weight)
+        return Spool.from_db(db_item)
+
+    if body.use_length is not None:
+        db_item = await spool.use_length(db, spool_id, body.use_length)
+        return Spool.from_db(db_item)
+
+    return JSONResponse(
+        status_code=400,
+        content={"message": "Either use_weight or use_length must be specified."},
+    )
diff --git a/spoolson/database/spool.py b/spoolson/database/spool.py
index c0fbfe5..2ed3b10 100644
--- a/spoolson/database/spool.py
+++ b/spoolson/database/spool.py
@@ -7,6 +7,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from spoolson.database import filament, models
 from spoolson.exceptions import ItemCreateError, ItemNotFoundError
+from spoolson.math import weight_from_length
 
 
 async def create(
@@ -41,9 +42,9 @@ async def create(
     return db_item
 
 
-async def get_by_id(db: AsyncSession, spool_id: int) -> models.Spool:
+async def get_by_id(db: AsyncSession, spool_id: int, with_for_update: Optional[bool] = None) -> models.Spool:
     """Get a spool object from the database by the unique ID."""
-    spool = await db.get(models.Spool, spool_id)
+    spool = await db.get(models.Spool, spool_id, with_for_update=with_for_update)  # type: ignore  # noqa: PGH003
     if spool is None:
         raise ItemNotFoundError(f"No spool with ID {spool_id} found.")
     return spool
@@ -70,3 +71,50 @@ async def delete(db: AsyncSession, spool_id: int) -> None:
     """Delete a spool object."""
     spool = await get_by_id(db, spool_id)
     await db.delete(spool)
+
+
+# TODO: Make unit tests for race conditions on these
+async def use_weight(db: AsyncSession, spool_id: int, weight: float) -> models.Spool:
+    """Reduce the weight of a spool.
+
+    Does nothing if the spool is empty.
+
+    Args:
+        db (AsyncSession): Database session
+        spool_id (int): Spool ID
+        weight (float): Weight loss in grams
+
+    Returns:
+        models.Spool: Updated spool object
+    """
+    spool = await get_by_id(db, spool_id, with_for_update=True)
+    spool.weight -= min(spool.weight, weight)
+    await db.flush()
+    return spool
+
+
+async def use_length(db: AsyncSession, spool_id: int, length: float) -> models.Spool:
+    """Reduce the weight of a spool by using a length of filament.
+
+    Does nothing if the spool is empty.
+
+    Args:
+        db (AsyncSession): Database session
+        spool_id (int): Spool ID
+        length (float): Length of filament to reduce by
+
+    Returns:
+        models.Spool: Updated spool object
+    """
+    spool = await get_by_id(db, spool_id, with_for_update=True)
+
+    filament = spool.filament
+
+    weight = weight_from_length(
+        length=length,
+        radius=filament.diameter / 2,
+        density=filament.density,
+    )
+    spool.weight -= min(spool.weight, weight)
+    await db.flush()
+    return spool
diff --git a/spoolson/math.py b/spoolson/math.py
new file mode 100644
index 0000000..87a7719
--- /dev/null
+++ b/spoolson/math.py
@@ -0,0 +1,19 @@
+"""Various math-related functions."""
+
+import math
+
+
+def weight_from_length(*, length: float, radius: float, density: float) -> float:
+    """Calculate the weight of a piece of filament.
+
+    Args:
+        length (float): Filament length in mm
+        radius (float): Filament radius in mm
+        density (float): Density of filament material in g/cm3
+
+    Returns:
+        float: Weight in g
+    """
+    volume_mm3 = length * math.pi * radius * radius
+    volume_cm3 = volume_mm3 / 1000
+    return density * volume_cm3