hopefully silence stupid pandera warning

1 files changed, 4 insertions, 5 deletions

diff --git a/generalresearch/models/thl/survey/task_collection.py b/generalresearch/models/thl/survey/task_collection.py
index b09cdfb..b80166f 100644
--- a/generalresearch/models/thl/survey/task_collection.py
+++ b/generalresearch/models/thl/survey/task_collection.py
@@ -5,8 +5,7 @@ import json
 import logging
 
 import pandas as pd
-from pandera import DataFrameSchema
-from pandera.errors import SchemaErrors
+import pandera.pandas as pa
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from generalresearch.models.thl.survey import MarketplaceTask
@@ -29,14 +28,14 @@ class TaskCollection(BaseModel):
     df: pd.DataFrame = Field(default_factory=pd.DataFrame)
 
     # overload this with the correct schema!
-    _schema: DataFrameSchema
+    _schema: pa.DataFrameSchema
 
     @model_validator(mode="after")
     def handle_df(self):
         df = self.to_df()
         try:
             df = self._schema.validate(df, lazy=True)
-        except SchemaErrors as exc:
+        except pa.errors.SchemaErrors as exc:
             idx = exc.failure_cases["index"]
             if len(idx) >= len(df) * 0.10:
                 raise exc
@@ -50,7 +49,7 @@ class TaskCollection(BaseModel):
     def to_df(self) -> pd.DataFrame: ...
 
 
-def create_empty_df_from_schema(schema: DataFrameSchema) -> pd.DataFrame:
+def create_empty_df_from_schema(schema: pa.DataFrameSchema) -> pd.DataFrame:
     # Create an empty df from the schema. We have to do this or else a plain empty df
     #   will fail validating non-nullable columns b/c they don't have a default.
     schema = copy.deepcopy(schema)