Refactor type hints and improve parameter handling in capacity models; enhance CsvList class methods

Author: btoron

ofsc/capacity.py

@@ -1,14 +1,18 @@
-from typing import Optional, Union, get_origin, get_args
+from typing import Optional, Union, get_args, get_origin
 from urllib.parse import urljoin
-import inspect
 
 import requests
 from pydantic import BaseModel
 
 from ofsc.models import CsvList, OFSApi
 
 from .common import OBJ_RESPONSE, wrap_return
-from .models import CapacityRequest, GetCapacityResponse, GetQuotaRequest, GetQuotaResponse
+from .models import (
+    CapacityRequest,
+    GetCapacityResponse,
+    GetQuotaRequest,
+    GetQuotaResponse,
+)
 
 
 def _convert_model_to_api_params(model: BaseModel) -> dict:
@@ -64,7 +68,7 @@ class OFSCapacity(OFSApi):
     @wrap_return(response_type=OBJ_RESPONSE, model=GetCapacityResponse)
     def getAvailableCapacity(self, 
                              dates: Union[list[str], CsvList, str],
-                             areas: Union[list[str], CsvList, str],
+                             areas: Optional[Union[list[str], CsvList, str]] = None,
                              categories: Optional[Union[list[str], CsvList, str]] = None,
                              aggregateResults: Optional[bool] = None,
                              availableTimeIntervals: str = "all",
@@ -164,3 +168,8 @@ def getQuota(self,
             params=params,
         )
         return response
+            url,
+            headers=self.headers,
+            params=params,
+        )
+        return response

ofsc/models.py

@@ -27,37 +27,37 @@
 
 class CsvList(BaseModel):
     """Auxiliary model to represent a list of strings as comma-separated values"""
-    
+
     value: str = ""
-    
+
     @classmethod
-    def from_list(cls, string_list: List[str]) -> 'CsvList':
+    def from_list(cls, string_list: List[str]) -> "CsvList":
         """Create CsvList from a list of strings
-        
+
         Args:
             string_list: List of strings to convert to CSV format
-            
+
         Returns:
             CsvList instance with comma-separated values
         """
         if not string_list:
             return cls(value="")
         return cls(value=",".join(string_list))
-    
+
     def to_list(self) -> List[str]:
         """Convert CsvList to a list of strings
-        
+
         Returns:
             List of strings split by commas, empty list if value is empty
         """
         if not self.value or self.value.strip() == "":
             return []
         return [item.strip() for item in self.value.split(",") if item.strip()]
-    
+
     def __str__(self) -> str:
         """String representation returns the CSV value"""
         return self.value
-    
+
     def __repr__(self) -> str:
         """Representation shows both CSV and list format"""
         return f"CsvList(value='{self.value}', list={self.to_list()})"
@@ -989,18 +989,19 @@ class DailyExtractFiles(BaseModel):
 
 class CapacityRequest(BaseModel):
     """Request model for capacity queries with CsvList support for string arrays
-    
+
     Accepts list[str], CsvList, or str for string array parameters but converts internally to CsvList
     """
+
     aggregateResults: Optional[bool] = None
-    areas: CsvList
+    areas: Optional[CsvList] = None
     availableTimeIntervals: str = "all"
     calendarTimeIntervals: str = "all"
     categories: Optional[CsvList] = None
     dates: CsvList
     fields: Optional[CsvList] = None
-    
-    @field_validator('areas', 'categories', 'dates', 'fields', mode='before')
+
+    @field_validator("areas", "categories", "dates", "fields", mode="before")
     @classmethod
     def convert_to_csvlist(cls, v):
         """Convert list[str], CsvList, str, or dict to CsvList"""
@@ -1013,24 +1014,26 @@ def convert_to_csvlist(cls, v):
         elif isinstance(v, str):
             # Handle string input as CSV
             return CsvList(value=v)
-        elif isinstance(v, dict) and 'value' in v:
+        elif isinstance(v, dict) and "value" in v:
             # Handle dict from JSON deserialization
-            return CsvList(value=v['value'])
+            return CsvList(value=v["value"])
         else:
-            raise ValueError(f"Expected list[str], CsvList, str, dict with 'value' key, or None, got {type(v)}")
-    
+            raise ValueError(
+                f"Expected list[str], CsvList, str, dict with 'value' key, or None, got {type(v)}"
+            )
+
     def get_areas_list(self) -> List[str]:
         """Get areas as a list of strings"""
         return self.areas.to_list()
-    
+
     def get_categories_list(self) -> List[str]:
         """Get categories as a list of strings"""
         return self.categories.to_list() if self.categories is not None else []
-    
+
     def get_dates_list(self) -> List[str]:
         """Get dates as a list of strings"""
         return self.dates.to_list()
-    
+
     def get_fields_list(self) -> List[str]:
         """Get fields as a list of strings"""
         return self.fields.to_list() if self.fields is not None else []
@@ -1076,7 +1079,7 @@ class GetCapacityResponse(BaseModel):
 
 class QuotaAreaItem(BaseModel):
     """Model for quota area items with quota-specific fields"""
-    
+
     label: Optional[str] = None
     name: Optional[str] = None
     maxAvailable: Optional[int] = None
@@ -1092,7 +1095,7 @@ class QuotaAreaItem(BaseModel):
 
 class QuotaResponseItem(BaseModel):
     """Model for individual quota response item by date"""
-    
+
     date: str
     areas: List[QuotaAreaItem] = []
 
@@ -1105,9 +1108,10 @@ class GetQuotaResponse(BaseModel):
 
 class GetQuotaRequest(BaseModel):
     """Request model for quota queries with comprehensive parameters
-    
+
     Accepts list[str] or CsvList for string array parameters but converts internally to CsvList
     """
+
     aggregateResults: Optional[bool] = None
     areas: Optional[CsvList] = None
     categories: Optional[CsvList] = None
@@ -1116,8 +1120,8 @@ class GetQuotaRequest(BaseModel):
     intervalLevel: Optional[bool] = None
     returnStatuses: Optional[bool] = None
     timeSlotLevel: Optional[bool] = None
-    
-    @field_validator('areas', 'categories', 'dates', mode='before')
+
+    @field_validator("areas", "categories", "dates", mode="before")
     @classmethod
     def convert_to_csvlist(cls, v):
         """Convert list[str], CsvList, str, or dict to CsvList"""
@@ -1130,20 +1134,22 @@ def convert_to_csvlist(cls, v):
         elif isinstance(v, str):
             # Handle string input as CSV
             return CsvList(value=v)
-        elif isinstance(v, dict) and 'value' in v:
+        elif isinstance(v, dict) and "value" in v:
             # Handle dict from JSON deserialization
-            return CsvList(value=v['value'])
+            return CsvList(value=v["value"])
         else:
-            raise ValueError(f"Expected list[str], CsvList, str, dict with 'value' key, or None, got {type(v)}")
-    
+            raise ValueError(
+                f"Expected list[str], CsvList, str, dict with 'value' key, or None, got {type(v)}"
+            )
+
     def get_areas_list(self) -> List[str]:
         """Get areas as a list of strings"""
         return self.areas.to_list() if self.areas is not None else []
-    
+
     def get_categories_list(self) -> List[str]:
         """Get categories as a list of strings"""
         return self.categories.to_list() if self.categories is not None else []
-    
+
     def get_dates_list(self) -> List[str]:
         """Get dates as a list of strings"""
         return self.dates.to_list()