fix: update pydantic configuration for models

msk-mind · Nov 15, 2023 · 965dd6c · 965dd6c
1 parent afc13cc
commit 965dd6c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/src/luna/common/models.py b/src/luna/common/models.py
@@ -3,7 +3,7 @@
 import pandera as pa
 from pandera.engines.pandas_engine import PydanticModel
 from pandera.typing import Series
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class Slide(BaseModel):
@@ -19,9 +19,9 @@ class Slide(BaseModel):
     channel1_R: Optional[float] = None
     channel1_G: Optional[float] = None
     channel1_B: Optional[float] = None
+    properties: Optional[dict] = None
 
-    class Config:
-        extra = "allow"
+    model_config = ConfigDict(extra="allow")
 
 
 class ShapeFeaturesSchema(pa.DataFrameModel):
@@ -47,8 +47,7 @@ class Tile(BaseModel):
     tile_size: int
     tile_units: str
 
-    class Config:
-        extra = "allow"
+    model_config = ConfigDict(extra="allow")
 
 
 class StoredTile(Tile):

diff --git a/src/luna/pathology/cli/slide_etl.py b/src/luna/pathology/cli/slide_etl.py
@@ -13,7 +13,7 @@
 from tiffslide import TiffSlide
 
 from luna.common.dask import configure_dask_client, get_or_create_dask_client
-from luna.common.models import Slide
+from luna.common.models import Slide, SlideSchema
 from luna.common.utils import apply_csv_filter, get_config, timed
 from luna.pathology.common.utils import (
     get_downscaled_thumbnail,
@@ -249,7 +249,7 @@ def __slide_etl(
             output_urlpath,
         )
 
-    return pd.json_normalize(slide.__dict__)
+    return DataFrame[SlideSchema](pd.json_normalize(slide.__dict__))
 
 
 def fire_cli():