From 8ab2f4cbf1960290c285517be737e422ca5552a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maurice=20Hern=C3=A9?= <maurice.herne@rwth-aachen.de>
Date: Mon, 28 Oct 2024 14:59:03 +0100
Subject: [PATCH] Better way to do the attachments fix
---
.../notebook_models/notebook_cell_model.py | 28 +++++++++++++++++--
1 file changed, 26 insertions(+), 2 deletions(-)
diff --git a/src/Notebook_Processor/notebook_models/notebook_cell_model.py b/src/Notebook_Processor/notebook_models/notebook_cell_model.py
index 5f6c4d1..f1303ac 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[
--
GitLab