switch to vanilla enum and fix bug in page header processor and json …

…dump
VikParuchuri · Dec 3, 2024 · fe39d8b · fe39d8b
1 parent 11349b2
commit fe39d8b
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/marker/processors/page_header.py b/marker/processors/page_header.py
@@ -8,7 +8,7 @@ class PageHeaderProcessor(BaseProcessor):
     """
     A processor for moving PageHeaders to the top
     """
-    block_types = (BlockTypes.PageHeader)
+    block_types = (BlockTypes.PageHeader,)
 
     def __call__(self, document: Document):
         for page in document.pages:

diff --git a/marker/renderers/json.py b/marker/renderers/json.py
@@ -23,7 +23,7 @@ class JSONBlockOutput(BaseModel):
 
 class JSONOutput(BaseModel):
     children: List[JSONBlockOutput]
-    block_type: BlockTypes = BlockTypes.Document
+    block_type: str = str(BlockTypes.Document)
     metadata: dict
 
 

diff --git a/marker/schema/__init__.py b/marker/schema/__init__.py
@@ -1,7 +1,7 @@
-from enum import auto, StrEnum
+from enum import auto, Enum
 
 
-class BlockTypes(StrEnum):
+class BlockTypes(str, Enum):
     Line = auto()
     Span = auto()
     FigureGroup = auto()