Pass format

Author: nikhilwoodruff

scripts/create_tables.py

@@ -11,27 +11,13 @@
 # Add src to path
 sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
 
-from sqlmodel import SQLModel, create_engine, text
 from rich.console import Console
+from sqlmodel import SQLModel, create_engine
 
 from policyengine_api.config.settings import settings
 from policyengine_api.services.storage import get_service_role_client
 
 # Import all models to register them with SQLModel.metadata
-from policyengine_api.models import (
-    Dataset,
-    DatasetVersion,
-    Dynamic,
-    Parameter,
-    ParameterValue,
-    Policy,
-    Simulation,
-    TaxBenefitModel,
-    TaxBenefitModelVersion,
-    Variable,
-    AggregateOutput,
-    ChangeAggregate,
-)
 
 console = Console()
 
@@ -126,4 +112,3 @@ def apply_migrations(engine):
     engine = create_tables()
     apply_migrations(engine)
     console.print("\n[bold green]✓ Database setup complete!")
-

scripts/seed.py

@@ -1,40 +1,51 @@
-"""Seed database with UK and US tax-benefit models, variables, parameters, and datasets."""
+"""Seed database with UK and US models, variables, parameters, datasets."""
 
-import sys
-from pathlib import Path
 import logging
+import sys
 import warnings
+from pathlib import Path
+
 import logfire
 
 # Disable all SQLAlchemy and database logging BEFORE any imports
 logging.basicConfig(level=logging.ERROR)
-logging.getLogger('sqlalchemy').setLevel(logging.ERROR)
-warnings.filterwarnings('ignore')
+logging.getLogger("sqlalchemy").setLevel(logging.ERROR)
+warnings.filterwarnings("ignore")
 
 # Add src to path
 sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
 
-from policyengine.tax_benefit_models.uk import uk_latest
-from policyengine.tax_benefit_models.uk.datasets import ensure_datasets as ensure_uk_datasets
-from policyengine.tax_benefit_models.us import us_latest
-from policyengine.tax_benefit_models.us.datasets import ensure_datasets as ensure_us_datasets
-from policyengine_api.models import (
+from policyengine.tax_benefit_models.uk import uk_latest  # noqa: E402
+from policyengine.tax_benefit_models.uk.datasets import (  # noqa: E402
+    ensure_datasets as ensure_uk_datasets,
+)
+from policyengine.tax_benefit_models.us import us_latest  # noqa: E402
+from policyengine.tax_benefit_models.us.datasets import (  # noqa: E402
+    ensure_datasets as ensure_us_datasets,
+)
+from rich.console import Console  # noqa: E402
+from rich.progress import track  # noqa: E402
+from sqlmodel import Session, create_engine, select  # noqa: E402
+
+from policyengine_api.config.settings import settings  # noqa: E402
+from policyengine_api.models import (  # noqa: E402
+    Dataset,
+    Parameter,
+    ParameterValue,
     TaxBenefitModel,
     TaxBenefitModelVersion,
     Variable,
-    Parameter,
-    ParameterValue,
-    Dataset,
 )
-from policyengine_api.services.storage import upload_dataset_for_seeding
-from sqlmodel import Session, create_engine, select
-from policyengine_api.config.settings import settings
-from rich.console import Console
-from rich.progress import track
+from policyengine_api.services.storage import (  # noqa: E402
+    upload_dataset_for_seeding,
+)
 
 # Configure logfire
 if settings.logfire_token:
-    logfire.configure(token=settings.logfire_token, environment=settings.logfire_environment)
+    logfire.configure(
+        token=settings.logfire_token,
+        environment=settings.logfire_environment,
+    )
 
 console = Console()
 
@@ -49,12 +60,18 @@ def get_quiet_session():
 def seed_model(model_version, session) -> TaxBenefitModelVersion:
     """Seed a tax-benefit model with its variables and parameters."""
 
-    with logfire.span(f"seed_model", model=model_version.model.id, version=model_version.version):
+    with logfire.span(
+        "seed_model",
+        model=model_version.model.id,
+        version=model_version.version,
+    ):
         # Create or get the model
         console.print(f"[bold blue]Seeding {model_version.model.id}...")
 
         existing_model = session.exec(
-            select(TaxBenefitModel).where(TaxBenefitModel.name == model_version.model.id)
+            select(TaxBenefitModel).where(
+                TaxBenefitModel.name == model_version.model.id
+            )
         ).first()
 
         if existing_model:
@@ -79,7 +96,9 @@ def seed_model(model_version, session) -> TaxBenefitModelVersion:
         ).first()
 
         if existing_version:
-            console.print(f"  Model version {model_version.version} already exists, skipping")
+            console.print(
+                f"  Model version {model_version.version} already exists, skipping"
+            )
             return existing_version
 
         db_version = TaxBenefitModelVersion(
@@ -100,19 +119,26 @@ def seed_model(model_version, session) -> TaxBenefitModelVersion:
                     name=var.name,
                     entity=var.entity,
                     description=var.description or "",
-                    data_type=var.data_type.__name__ if hasattr(var.data_type, '__name__') else str(var.data_type),
+                    data_type=var.data_type.__name__
+                    if hasattr(var.data_type, "__name__")
+                    else str(var.data_type),
                     tax_benefit_model_version_id=db_version.id,
                 )
                 session.add(db_var)
 
             session.commit()
-            console.print(f"  [green]✓[/green] Added {len(model_version.variables)} variables")
+            console.print(
+                f"  [green]✓[/green] Added {len(model_version.variables)} variables"
+            )
 
         # Add parameters (creating a lookup for parameter values later)
         parameters_to_add = model_version.parameters
         if settings.limit_seed_parameters:
             parameters_to_add = model_version.parameters[:10_000]
-            console.print(f"  [yellow]Limiting to {len(parameters_to_add)} parameters (LIMIT_SEED_PARAMETERS=true)[/yellow]")
+            console.print(
+                f"  [yellow]Limiting to {len(parameters_to_add)} parameters "
+                f"(LIMIT_SEED_PARAMETERS=true)[/yellow]"
+            )
 
         with logfire.span("add_parameters", count=len(parameters_to_add)):
             console.print(f"  Adding {len(parameters_to_add)} parameters...")
@@ -121,9 +147,11 @@ def seed_model(model_version, session) -> TaxBenefitModelVersion:
             for param in track(parameters_to_add, description="Parameters"):
                 db_param = Parameter(
                     name=param.name,
-                    label=param.label if hasattr(param, 'label') else None,
+                    label=param.label if hasattr(param, "label") else None,
                     description=param.description or "",
-                    data_type=param.data_type.__name__ if hasattr(param.data_type, '__name__') else str(param.data_type),
+                    data_type=param.data_type.__name__
+                    if hasattr(param.data_type, "__name__")
+                    else str(param.data_type),
                     unit=param.unit,
                     tax_benefit_model_version_id=db_version.id,
                 )
@@ -132,22 +160,29 @@ def seed_model(model_version, session) -> TaxBenefitModelVersion:
                 session.refresh(db_param)
                 param_id_map[param.id] = db_param.id
 
-            console.print(f"  [green]✓[/green] Added {len(parameters_to_add)} parameters")
+            console.print(
+                f"  [green]✓[/green] Added {len(parameters_to_add)} parameters"
+            )
 
         # Add parameter values
         # Filter to only include values for parameters we actually added
         parameter_values_to_add = [
-            pv for pv in model_version.parameter_values
+            pv
+            for pv in model_version.parameter_values
             if pv.parameter.id in param_id_map
         ]
 
         with logfire.span("add_parameter_values", count=len(parameter_values_to_add)):
-            console.print(f"  Adding {len(parameter_values_to_add)} parameter values...")
+            console.print(
+                f"  Adding {len(parameter_values_to_add)} parameter values..."
+            )
             import math
 
             for pv in track(parameter_values_to_add, description="Parameter values"):
                 # Handle Infinity values - skip them as they can't be stored in JSON
-                if isinstance(pv.value, float) and (math.isinf(pv.value) or math.isnan(pv.value)):
+                if isinstance(pv.value, float) and (
+                    math.isinf(pv.value) or math.isnan(pv.value)
+                ):
                     continue
 
                 db_pv = ParameterValue(
@@ -159,7 +194,10 @@ def seed_model(model_version, session) -> TaxBenefitModelVersion:
                 session.add(db_pv)
 
             session.commit()
-            console.print(f"  [green]✓[/green] Added {len(parameter_values_to_add)} parameter values")
+            console.print(
+                f"  [green]✓[/green] Added {len(parameter_values_to_add)} "
+                f"parameter values"
+            )
 
         return db_version
 
@@ -171,34 +209,46 @@ def seed_datasets(session):
 
         # Get UK and US models
         uk_model = session.exec(
-            select(TaxBenefitModel).where(TaxBenefitModel.name.in_(["uk", "policyengine-uk"]))
+            select(TaxBenefitModel).where(
+                TaxBenefitModel.name.in_(["uk", "policyengine-uk"])
+            )
         ).first()
         us_model = session.exec(
-            select(TaxBenefitModel).where(TaxBenefitModel.name.in_(["us", "policyengine-us"]))
+            select(TaxBenefitModel).where(
+                TaxBenefitModel.name.in_(["us", "policyengine-us"])
+            )
         ).first()
 
         if not uk_model or not us_model:
-            console.print("[red]Error: UK or US model not found. Run seed_model first.[/red]")
+            console.print(
+                "[red]Error: UK or US model not found. Run seed_model first.[/red]"
+            )
             return
 
         # UK datasets
         console.print("  Creating UK datasets...")
         uk_datasets = ensure_uk_datasets()
 
         with logfire.span("seed_uk_datasets", count=len(uk_datasets)):
-            for _, pe_dataset in track(list(uk_datasets.items()), description="UK datasets"):
+            for _, pe_dataset in track(
+                list(uk_datasets.items()), description="UK datasets"
+            ):
                 # Check if dataset already exists
                 existing = session.exec(
                     select(Dataset).where(Dataset.name == pe_dataset.name)
                 ).first()
 
                 if existing:
-                    console.print(f"  Dataset {pe_dataset.name} already exists, skipping")
+                    console.print(
+                        f"  Dataset {pe_dataset.name} already exists, skipping"
+                    )
                     continue
 
                 # Upload to S3
                 object_name = upload_dataset_for_seeding(pe_dataset.filepath)
-                console.print(f"  Uploaded {pe_dataset.filepath} to S3 as {object_name}")
+                console.print(
+                    f"  Uploaded {pe_dataset.filepath} to S3 as {object_name}"
+                )
 
                 # Create database record
                 db_dataset = Dataset(
@@ -217,19 +267,25 @@ def seed_datasets(session):
         us_datasets = ensure_us_datasets()
 
         with logfire.span("seed_us_datasets", count=len(us_datasets)):
-            for _, pe_dataset in track(list(us_datasets.items()), description="US datasets"):
+            for _, pe_dataset in track(
+                list(us_datasets.items()), description="US datasets"
+            ):
                 # Check if dataset already exists
                 existing = session.exec(
                     select(Dataset).where(Dataset.name == pe_dataset.name)
                 ).first()
 
                 if existing:
-                    console.print(f"  Dataset {pe_dataset.name} already exists, skipping")
+                    console.print(
+                        f"  Dataset {pe_dataset.name} already exists, skipping"
+                    )
                     continue
 
                 # Upload to S3
                 object_name = upload_dataset_for_seeding(pe_dataset.filepath)
-                console.print(f"  Uploaded {pe_dataset.filepath} to S3 as {object_name}")
+                console.print(
+                    f"  Uploaded {pe_dataset.filepath} to S3 as {object_name}"
+                )
 
                 # Create database record
                 db_dataset = Dataset(
@@ -243,7 +299,9 @@ def seed_datasets(session):
                 session.commit()
                 console.print(f"  [green]✓[/green] Created dataset: {db_dataset.name}")
 
-        console.print(f"[green]✓[/green] Seeded {len(uk_datasets) + len(us_datasets)} datasets\n")
+        console.print(
+            f"[green]✓[/green] Seeded {len(uk_datasets) + len(us_datasets)} datasets\n"
+        )
 
 
 def main():

src/policyengine_api/config/settings.py

@@ -38,9 +38,13 @@ class Settings(BaseSettings):
     @property
     def database_url(self) -> str:
         """Get database URL from Supabase."""
-        return self.supabase_db_url or self.supabase_url.replace(
-            "http://", "postgresql://postgres:postgres@"
-        ).replace("https://", "postgresql://postgres:postgres@") + "/postgres"
+        return (
+            self.supabase_db_url
+            or self.supabase_url.replace(
+                "http://", "postgresql://postgres:postgres@"
+            ).replace("https://", "postgresql://postgres:postgres@")
+            + "/postgres"
+        )
 
 
 settings = Settings()

src/policyengine_api/main.py

@@ -1,11 +1,11 @@
 from contextlib import asynccontextmanager
 
+import logfire
 from fastapi import FastAPI
 from fastapi_cache import FastAPICache
 from fastapi_cache.backends.redis import RedisBackend
 from redis import asyncio as aioredis
 from rich.console import Console
-import logfire
 
 from policyengine_api.api import api_router
 from policyengine_api.config.settings import settings
@@ -14,12 +14,10 @@
 console = Console()
 
 # Configure Logfire
-from policyengine_api.config.settings import settings as app_settings
-
 logfire.configure(
     service_name="policyengine-api",
-    token=app_settings.logfire_token if app_settings.logfire_token else None,
-    environment=app_settings.logfire_environment,
+    token=settings.logfire_token if settings.logfire_token else None,
+    environment=settings.logfire_environment,
 )
 logfire.instrument_httpx()
 
@@ -32,7 +30,9 @@ async def lifespan(app: FastAPI):
     console.print("[bold green]Database initialized[/bold green]")
 
     console.print("[bold green]Initializing cache...[/bold green]")
-    redis = aioredis.from_url(settings.redis_url, encoding="utf8", decode_responses=True)
+    redis = aioredis.from_url(
+        settings.redis_url, encoding="utf8", decode_responses=True
+    )
     FastAPICache.init(RedisBackend(redis), prefix="fastapi-cache")
     console.print("[bold green]Cache initialized[/bold green]")
 

src/policyengine_api/models/dataset.py

@@ -5,8 +5,8 @@
 from sqlmodel import Field, Relationship, SQLModel
 
 if TYPE_CHECKING:
-    from .tax_benefit_model import TaxBenefitModel
     from .dataset_version import DatasetVersion
+    from .tax_benefit_model import TaxBenefitModel
 
 
 class DatasetBase(SQLModel):