document

PdfDocument

Bases: BaseModel

Represents a PDF document

Source code in docprompt/schema/document.py

class PdfDocument(BaseModel):  """  Represents a PDF document  """   name: str = Field(description="The name of the document")  file_bytes: bytes = Field(description="The bytes of the document", repr=False)  file_path: Optional[str] = None   password: Optional[SecretStr] = None   def __len__(self):  return self.num_pages   def __hash__(self):  return hash(self.document_hash)   @computed_field  @cached_property  def page_count(self) -> PositiveInt:  from docprompt.utils.util import get_page_count   return get_page_count(self.file_bytes)   @property  def num_pages(self):  return self.page_count   @property  def bytes_per_page(self):  return len(self.file_bytes) / self.num_pages   @computed_field  @cached_property  def document_hash(self) -> str:  from docprompt.utils.util import hash_from_bytes   return hash_from_bytes(self.file_bytes)   @field_serializer("file_bytes")  def serialize_file_bytes(self, v: bytes, _info):  compressed = gzip.compress(v)   return base64.b64encode(compressed).decode("utf-8")   @field_validator("file_bytes")  def validate_file_bytes(cls, v: bytes):  if not isinstance(v, bytes):  raise ValueError("File bytes must be bytes")   if len(v) == 0:  raise ValueError("File bytes must not be empty")   if filetype.guess_mime(v) == "text/plain":  v = base64.b64decode(v, validate=True)   if filetype.guess_mime(v) == "application/gzip":  v = gzip.decompress(v)   if filetype.guess_mime(v) != "application/pdf":  raise ValueError("File bytes must be a PDF")   return v   @classmethod  def from_path(cls, file_path: Union[PathLike, str]):  file_path = Path(file_path)   if not file_path.is_file():  raise ValueError(f"File path {file_path} is not a file")   file_bytes = file_path.read_bytes()   return cls(name=file_path.name, file_path=str(file_path), file_bytes=file_bytes)   @classmethod  def from_bytes(cls, file_bytes: bytes, name: Optional[str] = None):  if name is None:  name = f"PDF-{datetime.now().isoformat()}.pdf"   return cls(name=name, file_bytes=file_bytes)   def get_bytes(self) -> bytes:  return self.file_bytes # Deprecated   @property  def path(self):  return self.file_path   def get_page_render_size(  self, page_number: int, dpi: int = DEFAULT_DPI  ) -> Tuple[int, int]:  """  Returns the render size of a page in pixels  """  return get_page_render_size_from_bytes(self.get_bytes(), page_number, dpi=dpi)   def to_compressed_bytes(self, compression_kwargs: dict = {}) -> bytes:  """  Compresses the document using Ghostscript  """  with self.as_tempfile() as temp_path:  return compress_pdf_to_bytes(temp_path, **compression_kwargs)   def rasterize_page(  self,  page_number: int,  *,  dpi: int = DEFAULT_DPI,  downscale_size: Optional[Tuple[int, int]] = None,  resize_mode: ResizeModes = "thumbnail",  max_file_size_bytes: Optional[int] = None,  resize_aspect_ratios: Optional[Iterable[AspectRatioRule]] = None,  do_convert: bool = False,  image_convert_mode: str = "L",  do_quantize: bool = False,  quantize_color_count: int = 8,  return_mode: Literal["pil", "bytes"] = "bytes",  ):  """  Rasterizes a page of the document using Pdfium  """  if page_number <= 0 or page_number > self.num_pages:  raise ValueError(f"Page number must be between 0 and {self.num_pages}")   post_process_fn = None   if any(  (  downscale_size,  max_file_size_bytes,  resize_aspect_ratios,  do_convert,  do_quantize,  )  ):  post_process_fn = partial(  process_raster_image,  resize_width=downscale_size[0] if downscale_size else None,  resize_height=downscale_size[1] if downscale_size else None,  resize_mode=resize_mode,  resize_aspect_ratios=resize_aspect_ratios,  do_convert=do_convert,  image_convert_mode=image_convert_mode,  do_quantize=do_quantize,  quantize_color_count=quantize_color_count,  max_file_size_bytes=max_file_size_bytes,  )   rastered = rasterize_page_with_pdfium(  self.file_bytes,  page_number,  return_mode=return_mode,  post_process_fn=post_process_fn,  scale=(1 / 72) * dpi,  )   return rastered   def rasterize_page_to_data_uri(  self,  page_number: int,  *,  dpi: int = DEFAULT_DPI,  downscale_size: Optional[Tuple[int, int]] = None,  resize_mode: ResizeModes = "thumbnail",  max_file_size_bytes: Optional[int] = None,  resize_aspect_ratios: Optional[Iterable[AspectRatioRule]] = None,  do_convert: bool = False,  image_convert_mode: str = "L",  do_quantize: bool = False,  quantize_color_count: int = 8,  render_grayscale: bool = False,  ) -> str:  """  Rasterizes a page of the document using Pdfium and returns a data URI, which can  be embedded into HTML or passed to large language models  """  image_bytes = self.rasterize_page(  page_number,  dpi=dpi,  downscale_size=downscale_size,  do_convert=do_convert,  image_convert_mode=image_convert_mode,  do_quantize=do_quantize,  quantize_color_count=quantize_color_count,  resize_mode=resize_mode,  max_file_size_bytes=max_file_size_bytes,  resize_aspect_ratios=resize_aspect_ratios,  return_mode="bytes",  )  return f"data:image/png;base64,{base64.b64encode(image_bytes).decode('utf-8')}"   def rasterize_pdf(  self,  dpi: int = DEFAULT_DPI,  downscale_size: Optional[Tuple[int, int]] = None,  resize_mode: ResizeModes = "thumbnail",  max_file_size_bytes: Optional[int] = None,  resize_aspect_ratios: Optional[Iterable[AspectRatioRule]] = None,  do_convert: bool = False,  image_convert_mode: str = "L",  do_quantize: bool = False,  quantize_color_count: int = 8,  return_mode: Literal["pil", "bytes"] = "bytes",  render_grayscale: bool = False,  ) -> Dict[int, bytes]:  """  Rasterizes the entire document using Pdfium  """  result = {}   post_process_fn = None   if any(  (  downscale_size,  max_file_size_bytes,  resize_aspect_ratios,  do_convert,  do_quantize,  )  ):  post_process_fn = partial(  process_raster_image,  resize_width=downscale_size[0] if downscale_size else None,  resize_height=downscale_size[1] if downscale_size else None,  resize_mode=resize_mode,  resize_aspect_ratios=resize_aspect_ratios,  do_convert=do_convert,  image_convert_mode=image_convert_mode,  do_quantize=do_quantize,  quantize_color_count=quantize_color_count,  max_file_size_bytes=max_file_size_bytes,  )   for idx, rastered in enumerate(  rasterize_pdf_with_pdfium(  self.file_bytes,  scale=(1 / 72) * dpi,  grayscale=render_grayscale,  return_mode=return_mode,  post_process_fn=post_process_fn,  )  ):  result[idx + 1] = rastered   return result   def split(self, start: Optional[int] = None, stop: Optional[int] = None):  """  Splits a document into multiple documents  """  if start is None and stop is None:  raise ValueError("Must specify either start or stop")   start = start or 0   from docprompt.utils.splitter import split_pdf_to_bytes   split_bytes = split_pdf_to_bytes(  self.file_bytes, start_page=start, stop_page=stop  )   return Document.from_bytes(split_bytes, name=self.name)   def as_tempfile(self, **kwargs):  """  Returns a tempfile of the document  """   @contextmanager  def tempfile_context() -> Generator[str, None, None]:  tempfile_kwargs = {"mode": "wb", "delete": True, "suffix": ".pdf", **kwargs}   with tempfile.NamedTemporaryFile(**tempfile_kwargs) as f:  f.write(self.file_bytes)  f.flush()  yield f.name   return tempfile_context()   def write_to_path(self, path: Union[PathLike, str], **kwargs):  """  Writes the document to a path  """  path = Path(path)   if path.is_dir():  path = path / self.name   with path.open("wb") as f:  f.write(self.file_bytes) 

as_tempfile(**kwargs)

Returns a tempfile of the document

Source code in docprompt/schema/document.py

def as_tempfile(self, **kwargs):  """  Returns a tempfile of the document  """   @contextmanager  def tempfile_context() -> Generator[str, None, None]:  tempfile_kwargs = {"mode": "wb", "delete": True, "suffix": ".pdf", **kwargs}   with tempfile.NamedTemporaryFile(**tempfile_kwargs) as f:  f.write(self.file_bytes)  f.flush()  yield f.name   return tempfile_context() 

get_page_render_size(page_number, dpi=DEFAULT_DPI)

Returns the render size of a page in pixels

Source code in docprompt/schema/document.py

def get_page_render_size(  self, page_number: int, dpi: int = DEFAULT_DPI ) -> Tuple[int, int]:  """  Returns the render size of a page in pixels  """  return get_page_render_size_from_bytes(self.get_bytes(), page_number, dpi=dpi) 

rasterize_page(page_number, *, dpi=DEFAULT_DPI, downscale_size=None, resize_mode='thumbnail', max_file_size_bytes=None, resize_aspect_ratios=None, do_convert=False, image_convert_mode='L', do_quantize=False, quantize_color_count=8, return_mode='bytes')

Rasterizes a page of the document using Pdfium

Source code in docprompt/schema/document.py

def rasterize_page(  self,  page_number: int,  *,  dpi: int = DEFAULT_DPI,  downscale_size: Optional[Tuple[int, int]] = None,  resize_mode: ResizeModes = "thumbnail",  max_file_size_bytes: Optional[int] = None,  resize_aspect_ratios: Optional[Iterable[AspectRatioRule]] = None,  do_convert: bool = False,  image_convert_mode: str = "L",  do_quantize: bool = False,  quantize_color_count: int = 8,  return_mode: Literal["pil", "bytes"] = "bytes", ):  """  Rasterizes a page of the document using Pdfium  """  if page_number <= 0 or page_number > self.num_pages:  raise ValueError(f"Page number must be between 0 and {self.num_pages}")   post_process_fn = None   if any(  (  downscale_size,  max_file_size_bytes,  resize_aspect_ratios,  do_convert,  do_quantize,  )  ):  post_process_fn = partial(  process_raster_image,  resize_width=downscale_size[0] if downscale_size else None,  resize_height=downscale_size[1] if downscale_size else None,  resize_mode=resize_mode,  resize_aspect_ratios=resize_aspect_ratios,  do_convert=do_convert,  image_convert_mode=image_convert_mode,  do_quantize=do_quantize,  quantize_color_count=quantize_color_count,  max_file_size_bytes=max_file_size_bytes,  )   rastered = rasterize_page_with_pdfium(  self.file_bytes,  page_number,  return_mode=return_mode,  post_process_fn=post_process_fn,  scale=(1 / 72) * dpi,  )   return rastered 

rasterize_page_to_data_uri(page_number, *, dpi=DEFAULT_DPI, downscale_size=None, resize_mode='thumbnail', max_file_size_bytes=None, resize_aspect_ratios=None, do_convert=False, image_convert_mode='L', do_quantize=False, quantize_color_count=8, render_grayscale=False)

Rasterizes a page of the document using Pdfium and returns a data URI, which can be embedded into HTML or passed to large language models

Source code in docprompt/schema/document.py

def rasterize_page_to_data_uri(  self,  page_number: int,  *,  dpi: int = DEFAULT_DPI,  downscale_size: Optional[Tuple[int, int]] = None,  resize_mode: ResizeModes = "thumbnail",  max_file_size_bytes: Optional[int] = None,  resize_aspect_ratios: Optional[Iterable[AspectRatioRule]] = None,  do_convert: bool = False,  image_convert_mode: str = "L",  do_quantize: bool = False,  quantize_color_count: int = 8,  render_grayscale: bool = False, ) -> str:  """  Rasterizes a page of the document using Pdfium and returns a data URI, which can  be embedded into HTML or passed to large language models  """  image_bytes = self.rasterize_page(  page_number,  dpi=dpi,  downscale_size=downscale_size,  do_convert=do_convert,  image_convert_mode=image_convert_mode,  do_quantize=do_quantize,  quantize_color_count=quantize_color_count,  resize_mode=resize_mode,  max_file_size_bytes=max_file_size_bytes,  resize_aspect_ratios=resize_aspect_ratios,  return_mode="bytes",  )  return f"data:image/png;base64,{base64.b64encode(image_bytes).decode('utf-8')}" 

rasterize_pdf(dpi=DEFAULT_DPI, downscale_size=None, resize_mode='thumbnail', max_file_size_bytes=None, resize_aspect_ratios=None, do_convert=False, image_convert_mode='L', do_quantize=False, quantize_color_count=8, return_mode='bytes', render_grayscale=False)

Rasterizes the entire document using Pdfium

Source code in docprompt/schema/document.py

def rasterize_pdf(  self,  dpi: int = DEFAULT_DPI,  downscale_size: Optional[Tuple[int, int]] = None,  resize_mode: ResizeModes = "thumbnail",  max_file_size_bytes: Optional[int] = None,  resize_aspect_ratios: Optional[Iterable[AspectRatioRule]] = None,  do_convert: bool = False,  image_convert_mode: str = "L",  do_quantize: bool = False,  quantize_color_count: int = 8,  return_mode: Literal["pil", "bytes"] = "bytes",  render_grayscale: bool = False, ) -> Dict[int, bytes]:  """  Rasterizes the entire document using Pdfium  """  result = {}   post_process_fn = None   if any(  (  downscale_size,  max_file_size_bytes,  resize_aspect_ratios,  do_convert,  do_quantize,  )  ):  post_process_fn = partial(  process_raster_image,  resize_width=downscale_size[0] if downscale_size else None,  resize_height=downscale_size[1] if downscale_size else None,  resize_mode=resize_mode,  resize_aspect_ratios=resize_aspect_ratios,  do_convert=do_convert,  image_convert_mode=image_convert_mode,  do_quantize=do_quantize,  quantize_color_count=quantize_color_count,  max_file_size_bytes=max_file_size_bytes,  )   for idx, rastered in enumerate(  rasterize_pdf_with_pdfium(  self.file_bytes,  scale=(1 / 72) * dpi,  grayscale=render_grayscale,  return_mode=return_mode,  post_process_fn=post_process_fn,  )  ):  result[idx + 1] = rastered   return result 

split(start=None, stop=None)

Splits a document into multiple documents

Source code in docprompt/schema/document.py

def split(self, start: Optional[int] = None, stop: Optional[int] = None):  """  Splits a document into multiple documents  """  if start is None and stop is None:  raise ValueError("Must specify either start or stop")   start = start or 0   from docprompt.utils.splitter import split_pdf_to_bytes   split_bytes = split_pdf_to_bytes(  self.file_bytes, start_page=start, stop_page=stop  )   return Document.from_bytes(split_bytes, name=self.name) 

to_compressed_bytes(compression_kwargs={})

Compresses the document using Ghostscript

Source code in docprompt/schema/document.py

def to_compressed_bytes(self, compression_kwargs: dict = {}) -> bytes:  """  Compresses the document using Ghostscript  """  with self.as_tempfile() as temp_path:  return compress_pdf_to_bytes(temp_path, **compression_kwargs) 

write_to_path(path, **kwargs)

Writes the document to a path

Source code in docprompt/schema/document.py

def write_to_path(self, path: Union[PathLike, str], **kwargs):  """  Writes the document to a path  """  path = Path(path)   if path.is_dir():  path = path / self.name   with path.open("wb") as f:  f.write(self.file_bytes) 

get_page_render_size_from_bytes(file_bytes, page_number, dpi=DEFAULT_DPI)

Returns the render size of a page in pixels

Source code in docprompt/schema/document.py

def get_page_render_size_from_bytes(  file_bytes: bytes, page_number: int, dpi: int = DEFAULT_DPI ):  """  Returns the render size of a page in pixels  """   with get_pdfium_document(file_bytes) as pdf:  page = pdf.get_page(page_number)   mediabox = page.get_mediabox()   base_width = int(mediabox[2] - mediabox[0])  base_height = int(mediabox[3] - mediabox[1])   width = int(base_width * dpi / 72)  height = int(base_height * dpi / 72)   return width, height