Add helper functions for Pen

helpers/pen_helpers.py

@@ -1,67 +1,91 @@
 from typing import Optional
 from pydantic import BaseModel, validator
+from sqlalchemy.orm import Session
+from models import Pen
 
-class Pen(BaseModel):
+class PenCreate(BaseModel):
     name: str
     color: str
     brand: str
     price: int
-    stock: int
 
-    @validator('name')
+    @validator('name', 'color', 'brand')
-    def name_must_not_be_empty(cls, value):
+    def validate_string_not_empty(cls, value):
-        if not value.strip():
+        if not value or not value.strip():
-            raise ValueError('Name cannot be empty')
+            raise ValueError('Field cannot be empty or contain only whitespace characters')
-        return value
-
-    @validator('color')
-    def color_must_not_be_empty(cls, value):
-        if not value.strip():
-            raise ValueError('Color cannot be empty')
-        return value
-
-    @validator('brand')
-    def brand_must_not_be_empty(cls, value):
-        if not value.strip():
-            raise ValueError('Brand cannot be empty')
         return value
 
     @validator('price')
-    def price_must_be_positive(cls, value):
+    def validate_price(cls, value):
-        if value <= 0:
-            raise ValueError('Price must be a positive integer')
-        return value
-
-    @validator('stock')
-    def stock_must_be_non_negative(cls, value):
         if value < 0:
-            raise ValueError('Stock cannot be negative')
+            raise ValueError('Price cannot be negative')
         return value
 
-def create_pen(db, pen_data: Pen):
+def create_pen(db: Session, pen_data: PenCreate) -> Pen:
+    """
+    Create a new pen in the database.
+
+    Args:
+        db: Database session
+        pen_data: Pen data for creation
+
+    Returns:
+        Pen object
+    """
     pen = Pen(**pen_data.dict())
     db.add(pen)
     db.commit()
     db.refresh(pen)
     return pen
 
-def get_pen(db, pen_id: int) -> Optional[Pen]:
+def get_pen_by_name(db: Session, name: str) -> Optional[Pen]:
-    return db.query(Pen).filter(Pen.id == pen_id).first()
+    """
+    Get a pen by its name.
 
-def update_pen(db, pen_id: int, pen_data: Pen):
+    Args:
-    pen = get_pen(db, pen_id)
+        db: Database session
-    if not pen:
+        name: Name of the pen
-        return None
+
-    for field, value in pen_data.dict(exclude_unset=True).items():
+    Returns:
-        setattr(pen, field, value)
+        Pen object if found, None otherwise
+    """
+    return db.query(Pen).filter(Pen.name == name).first()
+
+def update_pen(db: Session, pen_id: int, pen_data: PenCreate) -> Optional[Pen]:
+    """
+    Update an existing pen in the database.
+
+    Args:
+        db: Database session
+        pen_id: ID of the pen to update
+        pen_data: Updated pen data
+
+    Returns:
+        Updated Pen object if successful, None otherwise
+    """
+    pen = db.query(Pen).filter(Pen.id == pen_id).first()
+    if pen:
+        for key, value in pen_data.dict(exclude_unset=True).items():
+            setattr(pen, key, value)
         db.commit()
         db.refresh(pen)
         return pen
-
-def delete_pen(db, pen_id: int):
-    pen = get_pen(db, pen_id)
-    if not pen:
     return None
+
+def delete_pen(db: Session, pen_id: int) -> bool:
+    """
+    Delete a pen from the database.
+
+    Args:
+        db: Database session
+        pen_id: ID of the pen to delete
+
+    Returns:
+        True if successful, False otherwise
+    """
+    pen = db.query(Pen).filter(Pen.id == pen_id).first()
+    if pen:
         db.delete(pen)
         db.commit()
-    return pen
+        return True
+    return False