diff --git a/src/Notebook_Processor/notebook_models/notebook_cell_model.py b/src/Notebook_Processor/notebook_models/notebook_cell_model.py index 5f6c4d119b1b3c77b79e8aee1e180ebaf1644f2a..f1303ac1bd69518caae8ccfc61c0bb35964f09ed 100644 --- a/src/Notebook_Processor/notebook_models/notebook_cell_model.py +++ b/src/Notebook_Processor/notebook_models/notebook_cell_model.py @@ -6,6 +6,7 @@ from typing import ( Annotated, Union, Literal, + Optional, Any, ) from uuid import uuid4 @@ -17,6 +18,7 @@ from pydantic import ( field_validator, ValidationInfo, ConfigDict, + model_serializer, ) from ..configuration import ( @@ -78,13 +80,35 @@ class BaseNotebookCellModel(BaseModel): raise PydanticUseDefault() return value + @model_serializer + def drop_optional_none_fields(self): + """ + Drops optional fields that are None. + + We should look for a more dynamic solution to this. + Currently those fields need to be hardcoded because there is no way to specifically + detect fields that are annotated as Optional because the matching only returns true + if you get it exactly (Optional != Optional[dict]). + + :return: The model with the optional fields that are None removed. + """ + optional_fields = [ + "attachments", + ] + drop = {key for key, value in self.__dict__.items() if value is None and key in optional_fields} + return { + key: value + for key, value in self.__dict__.items() + if key not in drop + } + class NotebookMarkdownCellModel(BaseNotebookCellModel): """ "Notebook markdown cell." """ cell_type: Literal["markdown"] - attachments: dict = Field(default={}, description="Media attachments (e.g. inline images), stored as mimebundle keyed by filename.") + attachments: Optional[dict] = Field(default=None, description="Media attachments (e.g. inline images), stored as mimebundle keyed by filename.") class NotebookCodeCellModel(BaseNotebookCellModel): @@ -101,7 +125,7 @@ class NotebookRawCellModel(BaseNotebookCellModel): "Notebook raw nbconvert cell." """ cell_type: Literal["raw"] - attachments: dict = Field(default={}, description="Media attachments (e.g. inline images), stored as mimebundle keyed by filename.") + attachments: Optional[dict] = Field(default=None, description="Media attachments (e.g. inline images), stored as mimebundle keyed by filename.") NotebookCellAnnotation = Annotated[