color identity in pull sheet

app/services/pull_sheet_service.py

@@ -1,4 +1,5 @@
 from typing import List, Dict
+import json
 import pandas as pd
 from datetime import datetime
 from pathlib import Path
@@ -9,7 +10,7 @@ import asyncio
 from app.schemas.file import FileInDB
 from app.services.base_service import BaseService
 from sqlalchemy.orm import Session
-
+from app.models.tcgplayer_products import TCGPlayerProduct, TCGPlayerGroup, MTGJSONSKU, MTGJSONCard
 
 logger = logging.getLogger(__name__)
 
@@ -48,7 +49,7 @@ class PullSheetService(BaseService):
         """
         try:
             # Read and process CSV data
-            items = await self._read_and_process_csv(file.path)
+            items = await self._read_and_process_csv(db, file.path)
             
             # Prepare template data
             template_data = {
@@ -79,8 +80,51 @@ class PullSheetService(BaseService):
         except Exception as e:
             logger.error(f"Error generating pull sheet PDF: {str(e)}")
             raise
+    
+    async def _get_color_identity(self, db: Session, row: pd.Series) -> str:
+        """Get color identity from a row.
+        
+        Args:
+            row: pandas Series
+        """
+        # get category id from set name
+        group_id = db.query(TCGPlayerGroup).filter(TCGPlayerGroup.name == row['Set']).first().group_id
+        # format number
+        number = str(int(row['Number'])) if 'Number' in row and pd.notna(row['Number']) and '/' not in str(row['Number']) else str(row['Number']) if 'Number' in row and pd.notna(row['Number']) and '/' in str(row['Number']) else ''
+        # get product info from category id
+        product_id = db.query(TCGPlayerProduct).filter(TCGPlayerProduct.group_id == group_id).filter(TCGPlayerProduct.name == row['Product Name']).filter(TCGPlayerProduct.ext_number == number).filter(TCGPlayerProduct.ext_rarity == row['Rarity']).first().tcgplayer_product_id
+        # get scryfall id from product id
+        mtgjson_id = db.query(MTGJSONSKU).filter(MTGJSONSKU.tcgplayer_product_id == product_id).first().mtgjson_uuid
+        scryfall_id = db.query(MTGJSONCard).filter(MTGJSONCard.mtgjson_uuid == mtgjson_id).first().scryfall_id
+        # get color identity from scryfall
+        scryfall_service = self.get_service('scryfall')
+        color_identity = await scryfall_service.get_color_identity(scryfall_id)
+        if color_identity is None:
+            return '?'
+        # color identity is str of json array, convert to human readable string of list
+        color_identity = [str(color) for color in color_identity]
+        # if color identity is empty, return C for colorless
+        if not color_identity:
+            return 'C'
+        # ensure order, W > U > B > R > G
+        color_identity = sorted(color_identity, key=lambda x: ['W', 'U', 'B', 'R', 'G'].index(x))
+        color_identity = ''.join(color_identity)
+        return color_identity
+
+    async def _update_row_color_identity(self, db: Session, row: pd.Series) -> pd.Series:
+        """Update color identity from a row.
+        
+        Args:
+            row: pandas Series
+        """
+        # get color identity from row
+        color_identity = await self._get_color_identity(db, row)
+        # update row with color identity
+        row['Color Identity'] = color_identity
+        return row
+
             
-    async def _read_and_process_csv(self, csv_path: str) -> List[Dict]:
+    async def _read_and_process_csv(self, db: Session, csv_path: str) -> List[Dict]:
         """Read and process CSV data using pandas.
         
         Args:
@@ -103,6 +147,15 @@ class PullSheetService(BaseService):
         
         # Sort by Set Release Date (descending) and then Product Name (ascending)
         df = df.sort_values(['Set Release Date', 'Set', 'Product Name'], ascending=[False, True, True])
+
+        # Process color identities for all rows
+        color_identities = []
+        for _, row in df.iterrows():
+            color_identity = await self._get_color_identity(db, row)
+            color_identities.append(color_identity)
+        
+        # Add color identity column to dataframe
+        df['Color Identity'] = color_identities
         
         # Convert to list of dictionaries
         items = []
@@ -113,7 +166,8 @@ class PullSheetService(BaseService):
                 'quantity': str(int(row['Quantity'])),  # Convert to string for template
                 'set': row['Set'],
                 'rarity': row['Rarity'],
-                'card_number': str(int(row['Number'])) if 'Number' in row and pd.notna(row['Number']) and '/' not in str(row['Number']) else str(row['Number']) if 'Number' in row and pd.notna(row['Number']) and '/' in str(row['Number']) else ''
+                'card_number': str(int(row['Number'])) if 'Number' in row and pd.notna(row['Number']) and '/' not in str(row['Number']) else str(row['Number']) if 'Number' in row and pd.notna(row['Number']) and '/' in str(row['Number']) else '',
+                'color_identity': row['Color Identity']
             })
                 
         return items