diff --git a/mindee/parsing/custom/list.py b/mindee/parsing/custom/list.py index 23b90bb2..b6aed712 100644 --- a/mindee/parsing/custom/list.py +++ b/mindee/parsing/custom/list.py @@ -28,7 +28,7 @@ class ListFieldV1: """Confidence score""" reconstructed: bool """Whether the field was reconstructed from other fields.""" - page_id: int + page_id: Optional[int] """The document page on which the information was found.""" values: List[ListFieldValueV1] """List of word values""" @@ -41,15 +41,19 @@ def __init__( ) -> None: self.values = [] self.reconstructed = reconstructed + + for value in raw_prediction["values"]: + self.values.append(ListFieldValueV1(value)) + if "page_id" in value: + page_id = value["page_id"] + if page_id is None: - self.page_id = raw_prediction["page_id"] + if "page_id" in raw_prediction: + self.page_id = raw_prediction["page_id"] else: self.page_id = page_id self.confidence = raw_prediction["confidence"] - for value in raw_prediction["values"]: - self.values.append(ListFieldValueV1(value)) - @property def contents_list(self) -> List[str]: """Return a List of the contents of all values."""