leaf_focus.pdf.xpdf

  1"""Text extraction from pdf using xpdf tools."""
  2
  3from __future__ import annotations
  4
  5import dataclasses
  6import json
  7import logging
  8import pathlib
  9import subprocess
 10
 11from datetime import datetime
 12
 13from beartype import beartype, typing
 14from defusedxml import ElementTree
 15
 16from leaf_focus import utils
 17from leaf_focus.pdf import model
 18
 19
 20logger = logging.getLogger(__name__)
 21
 22
 23@beartype
 24class XpdfProgram:
 25    """Interact with xpdf tools."""
 26
 27    OPTS_TEXT_ENCODING: tuple[str, str, str, str, str, str] = (
 28        "Latin1",
 29        "ASCII7",
 30        "Symbol",
 31        "ZapfDingbats",
 32        "UTF-8",
 33        "UCS-2",
 34    )
 35    OPTS_TEXT_LINE_ENDING: tuple[str, str, str] = ("unix", "dos", "mac")
 36    OPTS_IMAGE_ROTATION: tuple[int, int, int, int] = (0, 90, 180, 270)
 37    OPTS_IMAGE_FREETYPE: tuple[str, str] = ("yes", "no")
 38    OPTS_IMAGE_ANTI_ALIAS: tuple[str, str] = ("yes", "no")
 39    OPTS_IMAGE_VEC_ANTI_ALIAS: tuple[str, str] = ("yes", "no")
 40
 41    def __init__(self, directory: pathlib.Path) -> None:
 42        """Create a new xpdf program class to interact with xpdf tools.
 43
 44        Args:
 45            directory: The path to the directory containing xpdf tools.
 46        """
 47        self._directory = directory
 48
 49    def info(
 50        self,
 51        pdf_path: pathlib.Path,
 52        output_dir: pathlib.Path,
 53        xpdf_args: model.XpdfInfoArgs,
 54    ) -> model.XpdfInfoResult:
 55        """Get information from a pdf file.
 56
 57        Args:
 58            pdf_path: The path to the pdf file.
 59            output_dir: The directory to save pdf info file.
 60            xpdf_args: The program arguments.
 61
 62        Returns:
 63            The pdf file information.
 64        """
 65        # validation
 66        enc = xpdf_args.encoding
 67        utils.validate("text encoding", enc, self.OPTS_TEXT_ENCODING)
 68
 69        utils.validate_pages(xpdf_args.first_page, xpdf_args.last_page)
 70
 71        if not pdf_path.exists():
 72            msg = f"Pdf file not found '{pdf_path}'."
 73            raise utils.LeafFocusError(msg) from FileNotFoundError(pdf_path)
 74
 75        output_file = utils.output_root(pdf_path, "info", output_dir)
 76        output_file = output_file.with_suffix(".json")
 77
 78        if output_file.exists():
 79            logger.info("Loading existing pdf info file.")
 80            with pathlib.Path.open(output_file, encoding="utf-8") as info_file:
 81                data = json.load(info_file)
 82                data["creation_date"] = utils.parse_date(data.get("creation_date"))
 83                data["modification_date"] = utils.parse_date(
 84                    data.get("modification_date")
 85                )
 86                return model.XpdfInfoResult(**data)
 87
 88        logger.info("Extracting pdf info and saving to file.")
 89
 90        # build command
 91        exe_path = utils.select_exe(self._directory / "pdfinfo")
 92        cmd = [str(exe_path)]
 93
 94        cmd_args = self.build_cmd(xpdf_args)
 95
 96        cmd.extend(cmd_args)
 97        cmd.append(str(pdf_path.resolve()))
 98
 99        # execute program
100        result = subprocess.run(  # noqa: S603
101            cmd,
102            capture_output=True,
103            check=True,
104            timeout=30,
105            text=True,
106        )
107        lines = result.stdout.splitlines()
108
109        metadata_line_index, data = self.build_field_metadata(
110            pdf_path,
111            lines,
112        )
113
114        # metadata
115        if metadata_line_index is not None:
116            start = metadata_line_index + 1
117            metadata = "\n".join(lines[start:])
118            root = ElementTree.fromstring(metadata)
119            data["metadata"] = utils.xml_to_element(root).to_dict()
120
121        if output_dir and output_dir.exists():
122            logger.debug("Saving pdf info to '%s'.", output_file)
123            output_file.write_text(
124                json.dumps(data, indent=2, cls=utils.CustomJsonEncoder),
125            )
126
127        return model.XpdfInfoResult(**data)
128
129    def text(
130        self,
131        pdf_path: pathlib.Path,
132        output_path: pathlib.Path,
133        xpdf_args: model.XpdfTextArgs,
134    ) -> model.XpdfTextResult:
135        """Get the text from a pdf file.
136
137        Args:
138            pdf_path: The path to the pdf file.
139            output_path: The directory to save output files.
140            xpdf_args: The pdf program arguments.
141
142        Returns:
143            The result from running the text extraction program.
144        """
145        # validation
146        eol = xpdf_args.line_end_type
147        utils.validate("end of line", eol, self.OPTS_TEXT_LINE_ENDING)
148
149        utils.validate_pages(xpdf_args.first_page, xpdf_args.last_page)
150
151        if not pdf_path.exists():
152            msg = f"Pdf file not found '{pdf_path}'."
153            raise utils.LeafFocusError(msg) from FileNotFoundError(str(pdf_path))
154
155        # build command
156
157        cmd_args = self.build_cmd(xpdf_args)
158
159        output_file = utils.output_root(pdf_path, "output", output_path, cmd_args)
160        output_file = output_file.with_suffix(".txt")
161
162        # check if embedded text file already exists
163        if output_file.exists():
164            logger.info("Loading extracted embedded text from existing file.")
165            return model.XpdfTextResult(
166                stdout=[],
167                stderr=[],
168                output_path=output_file,
169            )
170
171        if logger.isEnabledFor(logging.DEBUG):
172            logger.debug("Did not find expected output file '%s'", output_file.name)
173            logger.debug("Listing items in '%s'", output_file.parent)
174            item_count = 0
175            for item in output_file.parent.iterdir():
176                item_count += 1
177                logger.debug("Found item '%s'", item)
178            logger.debug("Found %s items in dir.", item_count)
179
180        logger.info("Extracting pdf embedded text and saving to file.")
181
182        exe_path = utils.select_exe(self._directory / "pdftotext")
183
184        cmd = [str(exe_path)]
185
186        cmd.extend([*cmd_args, str(pdf_path), str(output_file)])
187
188        # execute program
189        result = subprocess.run(  # noqa: S603
190            cmd,
191            capture_output=True,
192            check=True,
193            timeout=30,
194            text=True,
195        )
196
197        logger.debug("Saving pdf embedded text to '%s'.", output_file)
198
199        return model.XpdfTextResult(
200            stdout=(result.stdout or "").splitlines(),
201            stderr=(result.stderr or "").splitlines(),
202            output_path=output_file,
203        )
204
205    def image(
206        self,
207        pdf_path: pathlib.Path,
208        output_path: pathlib.Path,
209        xpdf_args: model.XpdfImageArgs,
210    ) -> model.XpdfImageResult:
211        """Create images of pdf pages.
212
213        Args:
214            pdf_path: The path to the pdf file.
215            output_path: The directory to save output files.
216            xpdf_args: The program arguments.
217
218        Returns:
219            The  pdf file image info.
220        """
221        # validation
222        rot = xpdf_args.rotation
223        utils.validate("rotation", rot, self.OPTS_IMAGE_ROTATION)
224
225        free_type = xpdf_args.free_type
226        utils.validate("freetype", free_type, self.OPTS_IMAGE_FREETYPE)
227
228        anti_alias = xpdf_args.anti_aliasing
229        utils.validate("anti-aliasing", anti_alias, self.OPTS_IMAGE_ANTI_ALIAS)
230
231        anti_alias_vec = xpdf_args.anti_aliasing
232        utils.validate(
233            "vector anti-aliasing",
234            anti_alias_vec,
235            self.OPTS_IMAGE_VEC_ANTI_ALIAS,
236        )
237
238        utils.validate_pages(xpdf_args.first_page, xpdf_args.last_page)
239
240        if not pdf_path.exists():
241            msg = f"Pdf file not found '{pdf_path}'."
242            raise utils.LeafFocusError(msg) from FileNotFoundError(str(pdf_path))
243
244        logger.info("Saving each pdf page as an image.")
245
246        # build command
247        cmd_args = self.build_cmd(xpdf_args)
248
249        output_type = "page-image"
250
251        # don't include the page limits when building the output prefix
252        xpdf_args.first_page = None
253        xpdf_args.last_page = None
254        output_cmd_args = self.build_cmd(xpdf_args)
255        output_dir = utils.output_root(
256            pdf_path,
257            output_type,
258            output_path,
259            output_cmd_args,
260        )
261
262        for pdf_image_file in output_dir.parent.iterdir():
263            if not pdf_image_file.name.startswith(output_dir.name):
264                continue
265
266            logger.info("Found existing pdf images.")
267
268            output_files = self.find_images(output_dir)
269            return model.XpdfImageResult(
270                stdout=[],
271                stderr=[],
272                output_dir=output_dir,
273                output_files=output_files,
274            )
275
276        exe_path = utils.select_exe(self._directory / "pdftopng")
277        cmd = [str(exe_path)]
278
279        cmd.extend([*cmd_args, str(pdf_path), str(output_dir)])
280
281        # execute program
282        result = subprocess.run(  # noqa: S603
283            cmd,
284            capture_output=True,
285            check=True,
286            timeout=30,
287            text=True,
288        )
289
290        logger.debug("Created pdf page images using prefix '%s'.", output_dir)
291
292        output_files = self.find_images(output_dir)
293
294        return model.XpdfImageResult(
295            stdout=(result.stdout or "").splitlines(),
296            stderr=(result.stderr or "").splitlines(),
297            output_dir=output_dir,
298            output_files=output_files,
299        )
300
301    def build_cmd(self, tool_args: model.XpdfArgs) -> list[str]:
302        """Build the command arguments from a data class."""
303        arg_class = tool_args.__class__
304        cmd_args = []
305        for field in dataclasses.fields(arg_class):
306            name = field.name
307            value = getattr(tool_args, name)
308
309            field_default = field.default
310
311            # TODO: account for default_factory
312
313            # validate the arg config
314            cmd_key = field.metadata.get("leaf_focus", {}).get("cmd")
315            if not cmd_key:
316                msg = f"Args incorrectly configured: missing 'cmd' for '{name}'."
317                raise ValueError(msg)
318
319            cmd_type = field.metadata.get("leaf_focus", {}).get("cmd_type")
320            if not cmd_type:
321                msg = f"Args incorrectly configured: missing 'cmd_type' for '{name}'."
322                raise ValueError(msg)
323
324            # add the arg
325            if cmd_type == "bool":
326                if value is not None and value is not True and value is not False:
327                    msg = (
328                        f"Argument '{name}' must be None, True, or False, "
329                        f"not '{value}'."
330                    )
331                    raise ValueError(msg)
332
333                if value is True:
334                    cmd_args.extend([str(cmd_key)])
335
336            elif cmd_type == "single":
337                if (
338                    field_default is None
339                    and value is not None
340                    or field_default != value
341                ):
342                    cmd_args.extend([str(cmd_key), str(value)])
343                else:
344                    # no need to add cmd
345                    pass
346            else:
347                msg = (
348                    f"Argument '{name}' has unknown cmd_type '{cmd_type}'. "
349                    "Expected one of 'bool, single'."
350                )
351                raise ValueError(msg)
352
353        return cmd_args
354
355    def find_images(self, output_dir: pathlib.Path) -> list[pathlib.Path]:
356        """Find image files in a directory."""
357        stem_parts = 7
358        stem_digit_parts = 6
359        output_files = []
360        for file_path in output_dir.parent.iterdir():
361            if not file_path.is_file():
362                continue
363            if not file_path.name.startswith(output_dir.stem):
364                continue
365            if file_path.suffix != ".png":
366                continue
367            if len(file_path.stem) < stem_parts:
368                continue
369            if file_path.stem[-stem_parts] != "-":
370                continue
371            if not all(i.isdigit() for i in file_path.stem[-stem_digit_parts:]):
372                continue
373            output_files.append(file_path)
374
375        if not output_files:
376            logger.warning("No page images found.")
377
378        return output_files
379
380    def build_field_metadata(
381        self,
382        pdf_path: pathlib.Path,
383        lines: typing.Iterable[str],
384    ) -> tuple[int | None, dict[str, typing.Any]]:
385        """Build metadata for a field."""
386        fields_map = {
387            field.metadata.get("leaf_focus", {}).get("name"): field
388            for field in dataclasses.fields(model.XpdfInfoResult)
389        }
390        metadata_line_index: int | None = None
391
392        data: dict[str, typing.Any] = {i.name: None for i in fields_map.values()}
393        for index, line in enumerate(lines):
394            if line.startswith("Metadata:"):
395                metadata_line_index = index
396                break
397
398            value: typing.Any = None
399            key, value = line.split(":", maxsplit=1)
400            key = key.strip()
401
402            field = fields_map.get(key)
403            if not field:
404                msg = f"Unknown pdf info key '{key}' value '{value}' in '{pdf_path}'."
405                raise utils.LeafFocusError(msg)
406
407            data_key = field.name
408            if data.get(data_key) is not None:
409                msg = f"Duplicate pdf info key '{key}' in '{pdf_path}'."
410                raise utils.LeafFocusError(msg)
411
412            typing_arg = typing.get_args(field.type)
413            types_str = [str, "str", "str | None"]
414            types_bool = [bool, "bool", "bool | None"]
415            types_int = [int, "int", "int | None"]
416            types_datetime = [datetime, "datetime", "datetime | None"]
417
418            if field.type in types_str or str in typing_arg:
419                value = value.strip()
420            elif field.type in types_datetime or datetime in typing_arg:
421                value = utils.parse_date(value.strip())
422            elif field.type in types_bool or bool in typing_arg:
423                value = value.strip().lower() == "yes"
424            elif field.type in types_int or int in typing_arg:
425                if data_key == "file_size_bytes":
426                    value = value.replace(" bytes", "")
427                value = int(value.strip().lower())
428            else:
429                msg = f"Unknown key '{key}' type '{field.type}'"
430                raise ValueError(msg)
431
432            data[data_key] = value
433
434        return metadata_line_index, data