hi-pdf-parser 0.0.4__py3-none-any.whl

hi_pdf_parser/__main__.py

@@ -0,0 +1,200 @@
+# Copyright (C) 2025 ByteDance Inc
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import json
+import logging
+from collections.abc import Iterator
+from enum import StrEnum
+from pathlib import Path
+from typing import Annotated, Any
+
+import typer
+import uvicorn
+
+from .app import create_app
+from .config import PyMuPDFParserConfig
+from .datamodel import Block
+from .parser import PyMuPDFParser
+from .settings import UvicornSettings
+
+app = typer.Typer(name="hi-pdf-parser", help="PDF Parser CLI and server")
+
+
+def _validate_pdf_file(file_path: str) -> Path:
+    p = Path(file_path)
+    if not p.exists():
+        raise typer.BadParameter(f"File does not exist: {file_path}")
+    if not p.is_file():
+        raise typer.BadParameter(f"Path is not a file: {file_path}")
+    return p
+
+
+class OutputFormat(StrEnum):
+    json = "json"
+    text = "text"
+
+
+def _output_results(
+    blocks: list[Block],
+    metadata: dict,
+    output_file: str | None,
+    output_format: OutputFormat,
+) -> None:
+    def _text_lines() -> Iterator[str]:
+        prev_page = -1
+        for b in blocks:
+            page_num: int | None = None
+            if b.areas:
+                page_num = b.areas[0].page_num
+            if page_num is not None and page_num != prev_page:
+                yield f"-------------- Page {page_num} --------------"
+                prev_page = page_num
+            yield b.content
+
+    if output_format == OutputFormat.text:
+        if output_file:
+            Path(output_file).write_text("\n".join(_text_lines()), encoding="utf-8")
+            typer.echo(f"Results written to: {output_file}")
+        else:
+            for line in _text_lines():
+                typer.echo(line)
+        return
+
+    result = {"blocks": [b.model_dump() for b in blocks], "metadata": metadata}
+    if output_file:
+        Path(output_file).write_text(
+            json.dumps(result, ensure_ascii=False), encoding="utf-8"
+        )
+        typer.echo(f"Results written to: {output_file}")
+    else:
+        typer.echo(json.dumps(result, ensure_ascii=False))
+
+
+@app.command()
+def parse(
+    pdf_file: Annotated[str, typer.Argument(help="Path to the PDF file to parse")],
+    output: Annotated[
+        str | None,
+        typer.Option("-o", "--output", help="Output file path (default: stdout)"),
+    ] = None,
+    extract_images: Annotated[
+        bool,
+        typer.Option(
+            "--extract-images/--no-extract-images", help="Extract images from the PDF"
+        ),
+    ] = True,
+    extract_tables: Annotated[
+        bool,
+        typer.Option(
+            "--extract-tables/--no-extract-tables", help="Extract tables from the PDF"
+        ),
+    ] = True,
+    skip_header_footer: Annotated[
+        bool,
+        typer.Option(
+            "--skip-header-footer/--no-skip-header-footer",
+            help="Skip header and footer detection",
+        ),
+    ] = True,
+    max_pages: Annotated[
+        int | None,
+        typer.Option(
+            "--max-pages",
+            help="Maximum number of pages to process (default: all pages)",
+        ),
+    ] = None,
+    password: Annotated[
+        str | None,
+        typer.Option("--password", help="Password for encrypted PDF files"),
+    ] = None,
+    format: Annotated[
+        OutputFormat, typer.Option("--format", help="Output format")
+    ] = OutputFormat.text,
+    verbose: Annotated[
+        bool, typer.Option("-v", "--verbose", help="Enable verbose logging")
+    ] = False,
+) -> None:
+    if verbose:
+        logging.basicConfig(level=logging.INFO)
+    pdf_path = _validate_pdf_file(pdf_file)
+    cfg: dict[str, Any] = {
+        "extract_images": extract_images,
+        "extract_tables": extract_tables,
+        "skip_header_footer": skip_header_footer,
+    }
+    if max_pages is not None:
+        cfg["max_pages"] = max_pages
+    config = PyMuPDFParserConfig(**cfg)
+    parser = PyMuPDFParser(config)
+    try:
+        typer.echo(f"Parsing PDF: {pdf_file}", err=True)
+        blocks, metadata = parser.parse(
+            str(pdf_path),
+            extract_images=extract_images,
+            extract_tables=extract_tables,
+            password=password,
+        )
+        typer.echo(f"Extracted {len(blocks)} blocks", err=True)
+        _output_results(blocks, metadata, output, format)
+    except FileNotFoundError as e:
+        typer.echo(f"Error: {e}", err=True)
+        raise typer.Exit(1)
+    except PermissionError as e:
+        typer.echo(f"Error: {e}", err=True)
+        raise typer.Exit(1)
+    except Exception as e:
+        typer.echo(f"Error parsing PDF: {e}", err=True)
+        raise typer.Exit(1)
+
+
+@app.command()
+def serve(
+    host: Annotated[str | None, typer.Option("--host", help="Server host")] = None,
+    port: Annotated[int | None, typer.Option("--port", help="Server port")] = None,
+    reload: Annotated[
+        bool | None, typer.Option("--reload/--no-reload", help="Enable auto-reload")
+    ] = None,
+    workers: Annotated[
+        int | None, typer.Option("--workers", help="Number of worker processes")
+    ] = None,
+    root_path: Annotated[
+        str | None, typer.Option("--root-path", help="Root path for the app")
+    ] = None,
+    proxy_headers: Annotated[
+        bool | None,
+        typer.Option("--proxy-headers/--no-proxy-headers", help="Use proxy headers"),
+    ] = None,
+    timeout_keep_alive: Annotated[
+        int | None,
+        typer.Option("--timeout-keep-alive", help="Keep-alive timeout seconds"),
+    ] = None,
+) -> None:
+    settings = UvicornSettings()
+    uvicorn.run(
+        app=create_app,
+        factory=True,
+        host=host or settings.host,
+        port=port or settings.port,
+        reload=reload if reload is not None else settings.reload,
+        workers=workers or settings.workers,
+        root_path=root_path or settings.root_path,
+        proxy_headers=proxy_headers
+        if proxy_headers is not None
+        else settings.proxy_headers,
+        timeout_keep_alive=timeout_keep_alive or settings.timeout_keep_alive,
+    )
+
+
+if __name__ == "__main__":
+    app()

hi_pdf_parser/app.py

@@ -0,0 +1,67 @@
+# Copyright (C) 2025 ByteDance Inc
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import logging
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.gzip import GZipMiddleware
+
+from .config import PyMuPDFParserConfig
+from .datamodel import HealthCheckResponse, ParseRequest, ParseResponse
+from .parser import PyMuPDFParser
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(levelname)s:\t%(asctime)s - %(name)s - %(message)s",
+    datefmt="%H:%M:%S",
+)
+
+_logger = logging.getLogger(__name__)
+
+
+def create_app() -> FastAPI:
+    parser = PyMuPDFParser(PyMuPDFParserConfig())
+    app = FastAPI(
+        title="PDF Parser Serve",
+    )
+
+    app.add_middleware(GZipMiddleware, minimum_size=5 * 1024, compresslevel=5)
+
+    @app.get("/health")
+    def health() -> HealthCheckResponse:
+        return HealthCheckResponse()
+
+    @app.post("/parse", response_model=ParseResponse)
+    def parse(request: ParseRequest):
+        try:
+            blocks, metadata = parser.parse(
+                request.file,
+                password=request.password,
+                extract_images=request.extract_images,
+                extract_tables=request.extract_tables,
+            )
+        except PermissionError as e:
+            raise HTTPException(
+                status_code=400, detail=f"Can not open encrypted file: {e}"
+            )
+        except FileNotFoundError:
+            raise HTTPException(
+                status_code=404, detail=f"File not found: {request.file}"
+            )
+        except Exception as e:
+            _logger.exception(f"Parse file fail: {request.file}")
+            raise HTTPException(status_code=500, detail=f"Parse extraction: {e}")
+        return ParseResponse(blocks=blocks, metadata=metadata)
+
+    return app

hi_pdf_parser/config.py

@@ -0,0 +1,24 @@
+# Copyright (C) 2025 ByteDance Inc
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from pydantic import BaseModel
+
+
+class PyMuPDFParserConfig(BaseModel):
+    """Configuration for PDF parser."""
+
+    extract_images: bool = True
+    extract_tables: bool = True
+    max_pages: int = 0  # 0 means no limit
+    skip_header_footer: bool = True

hi_pdf_parser/datamodel.py

@@ -0,0 +1,52 @@
+# Copyright (C) 2025 ByteDance Inc
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from enum import StrEnum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class HealthCheckResponse(BaseModel):
+    status: str = "ok"
+
+
+class ContentType(StrEnum):
+    image = "image"
+    text = "text"
+    table = "table"
+
+
+class BlockArea(BaseModel):
+    rect: tuple[float, float, float, float]
+    page_num: int
+
+
+class Block(BaseModel):
+    type: ContentType
+    areas: list[BlockArea]
+    content: str
+    font_sizes: list[float] = []
+
+
+class ParseRequest(BaseModel):
+    file: str = Field(description="file path", examples=["/path/my.pdf"])
+    password: str | None = Field(None, description="file password")
+    extract_images: bool = Field(True, description="enable image extraction")
+    extract_tables: bool = Field(True, description="enable table extraction")
+
+
+class ParseResponse(BaseModel):
+    blocks: list[Block]
+    metadata: dict[str, Any]