diff --git a/Readme.md b/Readme.md index 95dfbb3..d6a74ce 100644 --- a/Readme.md +++ b/Readme.md @@ -34,6 +34,8 @@ concurrent processing capabilities for PDF documents, reference strings, and pat - **JSON Output**: Convert TEI XML output to structured JSON format with CORD-19-like structure - **Markdown Output**: Convert TEI XML output to clean Markdown format with structured sections - **Type Hints**: Ships inline type annotations and a `py.typed` marker (PEP 561) for static type checking +- **Archive Streaming**: Process files directly from `.zip`/`.tar`/`.tar.gz` archives without fully decompressing them +- **S3 Streaming**: Read PDFs and zips straight from `s3://` (range-streamed, no full download) with the optional `[s3]` extra ## 📋 Prerequisites @@ -58,6 +60,9 @@ Choose one of the following installation methods: ```bash pip install grobid-client-python + +# to stream inputs directly from S3 (s3:// URIs), install the optional 's3' extra: +pip install "grobid-client-python[s3]" ``` ### Development Version @@ -167,8 +172,37 @@ grobid_client --server https://grobid.example.com --input ~/citations.txt proces # Force reprocessing with sentence segmentation and JSON output grobid_client --input ~/docs --force --segment_sentences --json processFulltextDocument + +# Process PDFs directly from a zip or tar.gz archive (streamed, not fully decompressed) +grobid_client --input ~/papers.zip --output ~/results processFulltextDocument +grobid_client --input ~/papers.tar.gz --output ~/results processFulltextDocument + +# --input also accepts glob patterns (quote them so the shell does not expand them) +grobid_client --input "~/papers/*.zip" --output ~/results processFulltextDocument # many archives +grobid_client --input "~/data/**/*.pdf" --output ~/results processFulltextDocument # PDFs in subdirectories ``` +> [!NOTE] +> `--input` accepts a directory, a single file, an **archive**, or a **glob pattern**: +> - **Archives** (`.zip`, `.tar`, `.tar.gz`/`.tgz`, `.tar.bz2`/`.tbz2`) are streamed: eligible entries are extracted in +> chunks of `batch_size` to a temporary directory, sent to GROBID, written to `--output`, and deleted before the next +> chunk. The archive is never fully decompressed, so disk usage stays bounded. If `--output` is omitted, results go to a +> directory named after the archive (e.g. `papers.zip` → `papers/`). +> - **Glob patterns** (`paper.zip`, `paper*.zip`, `**/paper*.zip`, `**/*.pdf`, …) are expanded with `**` recursion; each +> match is handled by type (archive → streamed, directory → recursed, file → processed). Quote the pattern so your shell +> passes it through to the client unexpanded. +> - **S3** (requires `pip install "grobid-client-python[s3]"`): pass an `s3://` object, prefix or glob. A remote zip is +> **range-streamed** (only its central directory and the entries are fetched — never the whole object); loose remote +> PDFs are fetched a batch at a time. Credentials use the standard AWS chain (env vars / `~/.aws` / IAM role). +> ```bash +> grobid_client --input "s3://my-bucket/papers/2021.zip" --output ~/out processFulltextDocument # one remote zip +> grobid_client --input "s3://my-bucket/pdfs/*.pdf" --output ~/out processFulltextDocument # loose PDFs +> grobid_client --input "s3://my-bucket/zips/" --output ~/out processFulltextDocument # every object under a prefix +> ``` +> +> A **manifest of paths** (local, glob or `s3://`, one per line, `#` comments allowed) can be processed together via +> `--input-list paths.txt` (combinable with `--input`). + ### Python Library #### Basic Usage diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index 8176103..69b86fc 100644 --- a/grobid_client/grobid_client.py +++ b/grobid_client/grobid_client.py @@ -19,6 +19,8 @@ import os import json import argparse +import fnmatch +import glob import time import concurrent.futures import ntpath @@ -26,13 +28,34 @@ import requests import pathlib import logging -from typing import Any, Optional, Tuple, Union +import shutil +import tarfile +import tempfile +import zipfile +from typing import Any, BinaryIO, Optional, Tuple, Union import copy from .format.TEI2LossyJSON import TEI2LossyJSONConverter from .client import ApiClient +def _default_file_mode() -> int: + """The mode open(..., 'w') would have produced, i.e. 0666 minus the umask. + + tempfile.mkstemp hardcodes 0600, so files written through it and renamed + into place would end up private -- unreadable to the group on shared + scratch, where the outputs of a cluster run usually have to be. Read the + umask once here, at import, because querying it means temporarily setting + it and that is not safe to do from worker threads. + """ + umask = os.umask(0o022) + os.umask(umask) + return 0o666 & ~umask + + +_DEFAULT_FILE_MODE = _default_file_mode() + + class ServerUnavailableException(Exception): """Exception raised when GROBID server is not available or not responding.""" @@ -47,6 +70,11 @@ class GrobidClient(ApiClient): # See https://github.com/grobidOrg/grobid-client-python/issues/54 CONSOLIDATE_CITATIONS_MIN_TIMEOUT = 120 + # Archive extensions that can be streamed entry-by-entry via --input instead + # of being fully decompressed first. Order matters: multi-dot suffixes must + # come before their single-dot prefixes when stripping (see _archive_stem). + ARCHIVE_EXTENSIONS = (".tar.gz", ".tar.bz2", ".tgz", ".tbz2", ".zip", ".tar") + # Default configuration values DEFAULT_CONFIG: dict = { 'grobid_server': 'http://localhost:8070', @@ -351,6 +379,55 @@ def _output_file_name( return str(filename) + def _write_atomic(self, filename: str, text: str) -> None: + """Write text to filename via a temp file in the same directory, then os.replace. + + A killed process must never leave a partial output behind. process_batch + decides a document is already done with os.path.isfile() alone, so a TEI + truncated by an OOM kill or a wall-clock timeout is indistinguishable + from a complete one and is skipped on every subsequent run -- the + corruption is permanent and silent. Writing to a temp file and renaming + means the destination either does not exist or is the whole document. + + The temp file goes in the DESTINATION directory, not TMPDIR: os.replace + is only atomic within a filesystem, and on a cluster TMPDIR is usually a + different mount. The "." prefix and ".tmp" suffix keep the temp file from + matching *.grobid.tei.xml or *_[0-9]*.txt, so output counting is + unaffected while a write is in flight. + + Residual risk: a SIGKILL between mkstemp and replace leaks a temp file. + That is visible and harmless, unlike a truncated TEI. + """ + dest = pathlib.Path(os.path.expanduser(filename)) + dest.parent.mkdir(parents=True, exist_ok=True) + # mkstemp names are unique, so concurrent writers from the + # ThreadPoolExecutor cannot collide on the temp path. + fd, tmp_path = tempfile.mkstemp(dir=str(dest.parent), prefix=".", suffix=".tmp") + try: + tmp_file = os.fdopen(fd, "w", encoding="utf8") + except BaseException: + # fdopen did not take ownership of fd, so we still have to close it. + # Past this point the file object owns it and closing it here too + # could close an unrelated descriptor that reused the number. + os.close(fd) + self._unlink_quietly(tmp_path) + raise + try: + with tmp_file: + tmp_file.write(text) + os.chmod(tmp_path, _DEFAULT_FILE_MODE) # mkstemp gives 0600 + os.replace(tmp_path, str(dest)) + except BaseException: + self._unlink_quietly(tmp_path) + raise + + @staticmethod + def _unlink_quietly(path: str) -> None: + try: + os.unlink(path) + except OSError: + pass + def ping(self) -> Tuple[bool, int]: """ Check the Grobid service. Returns True if the service is up. @@ -377,60 +454,663 @@ def process( json_output: bool = False, markdown_output: bool = False ) -> None: + if input_path is None: + self.logger.warning("No input path provided") + return + return self.process_paths( + service, [input_path], output=output, n=n, generate_ids=generate_ids, + consolidate_header=consolidate_header, consolidate_citations=consolidate_citations, + include_raw_citations=include_raw_citations, + include_raw_affiliations=include_raw_affiliations, + tei_coordinates=tei_coordinates, segment_sentences=segment_sentences, + force=force, verbose=verbose, flavor=flavor, + json_output=json_output, markdown_output=markdown_output, + ) + + def process_paths( + self, + service: str, + inputs: list, + output: Optional[str] = None, + n: int = 10, + generate_ids: bool = False, + consolidate_header: bool = True, + consolidate_citations: bool = False, + include_raw_citations: bool = False, + include_raw_affiliations: bool = False, + tei_coordinates: bool = False, + segment_sentences: bool = False, + force: bool = True, + verbose: bool = False, + flavor: Optional[str] = None, + json_output: bool = False, + markdown_output: bool = False + ) -> None: + """Process a list of inputs. + + Each input may be a local path, a shell glob (``**/*.pdf``), a directory, + a local archive, or an ``s3://`` object/prefix/glob. This backs both the + ``--input`` option (a single input) and ``--input-list`` (a manifest file + of paths). Results from all inputs are aggregated into one summary. + """ start_time = time.time() - batch_size_pdf = self.config["batch_size"] - # Warn if citation consolidation is requested with a short timeout: the - # consolidation step queries external services (e.g. CrossRef) and can - # be significantly slower, frequently resulting in HTTP 408 errors when - # the client-side timeout is too low. # See https://github.com/grobidOrg/grobid-client-python/issues/54 self._warn_on_consolidation_timeout(consolidate_citations) - # First pass: count all eligible files - all_input_files = [] - for (dirpath, dirnames, filenames) in os.walk(input_path): - for filename in filenames: - if filename.endswith(".pdf") or filename.endswith(".PDF") or \ - (service == 'processCitationList' and ( - filename.endswith(".txt") or filename.endswith(".TXT"))) or \ - (service == 'processCitationPatentST36' and ( - filename.endswith(".xml") or filename.endswith(".XML"))): - full_path = os.sep.join([dirpath, filename]) - all_input_files.append(full_path) - - # Log total files found - total_files = len(all_input_files) - if total_files == 0: - self.logger.warning(f"No eligible files found in {input_path}") + matched_paths = [] + for inp in inputs: + matched_paths.extend(self._resolve_input_paths(inp)) + if not matched_paths: + self.logger.warning(f"No files match input(s): {inputs}") + return + + # Partition into archives (streamed), remote loose files (s3) and local + # filesystem files (directories are expanded to their eligible files). + archive_paths = [] + remote_files = [] + fs_files = [] + for path in matched_paths: + if self._is_s3(path): + if self._looks_like_archive(path): + archive_paths.append(path) + elif self._is_eligible_input(self._s3_basename(path), service): + remote_files.append(path) + else: + self.logger.debug(f"Skipping s3 input (not an eligible file/archive): {path}") + elif self._is_archive(path): + archive_paths.append(path) + elif os.path.isdir(path): + fs_files.extend(self._collect_directory_files(path, service)) + elif os.path.isfile(path) and self._is_eligible_input(os.path.basename(path), service): + fs_files.append(path) + else: + self.logger.debug(f"Skipping input (not an eligible file/dir/archive): {path}") + + if not fs_files and not archive_paths and not remote_files: + self.logger.warning(f"No eligible files found in input(s): {inputs}") return - # Counters for processing statistics (initialize before early return) processed_files_count = 0 errors_files_count = 0 skipped_files_count = 0 + total_files = 0 + + # Local files gathered from directories and/or loose glob matches + if fs_files: + print(f"Found {len(fs_files)} local file(s) to process") + bp, be, bs = self._run_file_batches( + service, fs_files, self._common_base(fs_files), output, n, + generate_ids, consolidate_header, consolidate_citations, + include_raw_citations, include_raw_affiliations, tei_coordinates, + segment_sentences, force, verbose, flavor, json_output, markdown_output + ) + processed_files_count += bp + errors_files_count += be + skipped_files_count += bs + total_files += len(fs_files) + + # Loose remote (s3) files: streamed to a temp dir one chunk at a time + if remote_files: + rt, rp, re_count, rs = self._process_remote_files( + service, remote_files, output, n, + generate_ids, consolidate_header, consolidate_citations, + include_raw_citations, include_raw_affiliations, tei_coordinates, + segment_sentences, force, verbose, flavor, json_output, markdown_output + ) + processed_files_count += rp + errors_files_count += re_count + skipped_files_count += rs + total_files += rt + + # Archives (local or s3 zip) are streamed entry-by-entry per chunk + for archive_path in archive_paths: + at, ap, ae, as_count = self._process_archive_core( + service, archive_path, output, n, + generate_ids, consolidate_header, consolidate_citations, + include_raw_citations, include_raw_affiliations, tei_coordinates, + segment_sentences, force, verbose, flavor, json_output, markdown_output + ) + processed_files_count += ap + errors_files_count += ae + skipped_files_count += as_count + total_files += at + + if total_files == 0: + self.logger.warning(f"No eligible files found in input(s): {inputs}") + return + + runtime = time.time() - start_time + self._print_processing_summary( + processed_files_count, errors_files_count, skipped_files_count, total_files, runtime + ) + + def _resolve_input_paths(self, input_path: str) -> list: + """Resolve an input into a sorted list of concrete paths. + + Handles ``s3://`` URIs/prefixes/globs, shell-style glob patterns + (including the recursive ``**``) and ``~`` expansion. A plain local path + without glob metacharacters is returned as-is (so callers can still + handle a missing path themselves). + """ + if self._is_s3(input_path): + return self._resolve_s3_paths(input_path) + expanded = os.path.expanduser(input_path) + if glob.has_magic(expanded): + return sorted(glob.glob(expanded, recursive=True)) + return [expanded] + + def _collect_directory_files(self, directory: str, service: str) -> list: + """Recursively collect eligible input files from a directory.""" + files = [] + for path in sorted(pathlib.Path(directory).rglob('*')): + if path.is_file() and self._is_eligible_input(path.name, service): + files.append(str(path)) + return files + + def _common_base(self, files: list) -> str: + """Return a directory that is an ancestor of all given files. + + Used as ``input_path`` for output-name computation; only needs to be a + common ancestor so ``Path.relative_to`` does not fail. + """ + abs_files = [os.path.abspath(f) for f in files] + if len(abs_files) == 1: + return os.path.dirname(abs_files[0]) + try: + base = os.path.commonpath(abs_files) + except ValueError: + # e.g. paths on different drives (Windows); fall back to first parent + return os.path.dirname(abs_files[0]) + return base if os.path.isdir(base) else os.path.dirname(base) + + # ---- S3 support (optional 's3' extra: smart_open + boto3) ---- + + @staticmethod + def _is_s3(path: Any) -> bool: + """Return True if path is an s3:// URI.""" + return isinstance(path, str) and path.startswith("s3://") - print(f"Found {total_files} file(s) to process") - input_files = [] + @staticmethod + def _split_s3(uri: str) -> Tuple[str, str]: + """Split an s3://bucket/key URI into (bucket, key).""" + bucket, _, key = uri[len("s3://"):].partition("/") + return bucket, key - for input_file in all_input_files: - # Extract just the filename for verbose logging - filename = os.path.basename(input_file) + def _s3_basename(self, uri: str) -> str: + """Return the last path component of an s3:// key.""" + return self._split_s3(uri)[1].rsplit("/", 1)[-1] + + def _import_smart_open(self) -> Any: + try: + import smart_open # noqa: F401 + return smart_open + except ImportError as e: + raise ImportError( + "Reading from s3:// requires the optional 's3' extra. " + "Install it with: pip install grobid-client-python[s3]" + ) from e + + def _import_boto3(self) -> Any: + try: + import boto3 # noqa: F401 + return boto3 + except ImportError as e: + raise ImportError( + "Listing s3:// requires the optional 's3' extra. " + "Install it with: pip install grobid-client-python[s3]" + ) from e + + def _s3_open(self, uri: str) -> BinaryIO: + """Open an S3 object as a seekable binary stream (HTTP range-streamed). + + The returned stream lets zipfile read only the central directory and the + requested entries, so a remote zip is never fully downloaded. + """ + return self._import_smart_open().open(uri, "rb") + def _resolve_s3_paths(self, uri: str) -> list: + """Resolve an s3:// object/prefix/glob into a sorted list of object URIs. + + - ``s3://bucket/path/file.zip`` -> that single object + - ``s3://bucket/prefix/`` -> every object under the prefix + - ``s3://bucket/prefix/*.zip`` -> objects under the prefix matching the glob + """ + bucket, key = self._split_s3(uri) + if not bucket: + self.logger.warning(f"Invalid s3 uri: {uri}") + return [] + + pattern = None + if glob.has_magic(key): + magic = min(key.find(c) for c in "*?[" if c in key) + prefix = key[:magic] + pattern = key + elif key == "" or key.endswith("/"): + prefix = key + else: + return [uri] # a concrete object key + + s3 = self._import_boto3().client("s3") + keys = [] + for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=prefix): + for obj in page.get("Contents", []): + k = obj["Key"] + if k.endswith("/"): + continue + if pattern is None or fnmatch.fnmatch(k, pattern): + keys.append(k) + return [f"s3://{bucket}/{k}" for k in sorted(keys)] + + def _print_processing_summary( + self, + processed: int, + errors: int, + skipped: int, + total: int, + runtime: float + ) -> None: + """Print the final processing statistics (shared by all input modes).""" + docs_per_second = processed / runtime if runtime > 0 else 0 + seconds_per_doc = runtime / processed if processed > 0 else 0 + + print(f"Processing completed: {processed} out of {total} files processed") + print(f"Errors: {errors} out of {total} files processed") + if skipped > 0: + print(f"Skipped: {skipped} out of {total} files (already existed, use --force to reprocess)") + + print(f"⏱️ Total runtime: {runtime:.2f} seconds") + print(f"🚀 Speed: {docs_per_second:.2f} documents/second") + print(f" Throughput: {seconds_per_doc:.2f} seconds/document") + + def _run_file_batches( + self, + service: str, + input_files: list, + input_path: str, + output: Optional[str], + n: int, + generate_ids: bool, + consolidate_header: bool, + consolidate_citations: bool, + include_raw_citations: bool, + include_raw_affiliations: bool, + tei_coordinates: bool, + segment_sentences: bool, + force: bool, + verbose: bool, + flavor: Optional[str], + json_output: bool, + markdown_output: bool + ) -> Tuple[int, int, int]: + """Run process_batch over a list of files in chunks of batch_size. + + Returns the aggregated (processed, errors, skipped) counts. + """ + batch_size_pdf = self.config["batch_size"] + processed_files_count = 0 + errors_files_count = 0 + skipped_files_count = 0 + + batch = [] + for input_file in input_files: if verbose: try: - self.logger.info(f"Found file: {filename}") + self.logger.info(f"Found file: {os.path.basename(input_file)}") except UnicodeEncodeError: # may happen on linux see https://stackoverflow.com/questions/27366479/python-3-os-walk-file-paths-unicodeencodeerror-utf-8-codec-cant-encode-s - self.logger.warning(f"Could not log filename due to encoding issues") + self.logger.warning("Could not log filename due to encoding issues") + + batch.append(input_file) + + if len(batch) == batch_size_pdf: + batch_processed, batch_errors, batch_skipped = self.process_batch( + service, batch, input_path, output, n, generate_ids, + consolidate_header, consolidate_citations, include_raw_citations, + include_raw_affiliations, tei_coordinates, segment_sentences, + force, verbose, flavor, json_output, markdown_output + ) + processed_files_count += batch_processed + errors_files_count += batch_errors + skipped_files_count += batch_skipped + batch = [] + + if batch: + batch_processed, batch_errors, batch_skipped = self.process_batch( + service, batch, input_path, output, n, generate_ids, + consolidate_header, consolidate_citations, include_raw_citations, + include_raw_affiliations, tei_coordinates, segment_sentences, + force, verbose, flavor, json_output, markdown_output + ) + processed_files_count += batch_processed + errors_files_count += batch_errors + skipped_files_count += batch_skipped + + return processed_files_count, errors_files_count, skipped_files_count + + def _is_eligible_input(self, filename: str, service: str) -> bool: + """Return True if a file name is a valid input for the given service.""" + if filename.endswith(".pdf") or filename.endswith(".PDF"): + return True + if service == 'processCitationList' and ( + filename.endswith(".txt") or filename.endswith(".TXT")): + return True + if service == 'processCitationPatentST36' and ( + filename.endswith(".xml") or filename.endswith(".XML")): + return True + return False + + def _looks_like_archive(self, path: str) -> bool: + """Return True if the path/URI name has a known archive extension.""" + lower = path.lower() + return any(lower.endswith(ext) for ext in self.ARCHIVE_EXTENSIONS) + + def _is_archive(self, path: str) -> bool: + """Return True if path is an existing local zip/tar archive file.""" + return os.path.isfile(path) and self._looks_like_archive(path) + + def _archive_stem(self, path: str) -> str: + """Strip a known archive extension from path (e.g. docs.tar.gz -> docs).""" + lower = path.lower() + for ext in self.ARCHIVE_EXTENSIONS: + if lower.endswith(ext): + return path[:-len(ext)] + return os.path.splitext(path)[0] + + def _safe_member_path(self, dest_dir: str, arcname: str) -> Optional[str]: + """Resolve an archive entry name to a safe path under dest_dir. + + Leading slashes, drive letters and '..' components are stripped to + prevent path-traversal ("zip slip") outside of dest_dir. Returns None + if the entry name has no usable path component. + """ + normalized = arcname.replace("\\", "/") + parts = [p for p in normalized.split("/") if p not in ("", ".", "..")] + if not parts: + return None + return os.path.join(dest_dir, *parts) + + def _open_archive(self, archive_path: str) -> Tuple[str, Any, list]: + """Open a zip/tar archive and return (kind, handle, member_names). + + member_names contains only regular files (directories are skipped). + For s3:// zips the archive is range-streamed (not fully downloaded); the + underlying stream is stashed on the handle so the caller can close it. + + The handle is a ZipFile or a TarFile, which share no common interface + here: which one it is, is what the returned "kind" tag is for, and it is + the tag - not the type - that the callers dispatch on. + """ + archive: Any + if self._is_s3(archive_path): + if not archive_path.lower().endswith(".zip"): + raise ValueError( + f"Only .zip archives can be range-streamed over s3://: {archive_path}" + ) + stream = self._s3_open(archive_path) + archive = zipfile.ZipFile(stream) + archive._grobid_stream = stream # closed by _process_archive_core + names = [n for n in archive.namelist() if not n.endswith("/")] + return "zip", archive, names + + if archive_path.lower().endswith(".zip"): + archive = zipfile.ZipFile(archive_path) + names = [n for n in archive.namelist() if not n.endswith("/")] + return "zip", archive, names + + archive = tarfile.open(archive_path, "r:*") + names = [m.name for m in archive.getmembers() if m.isfile()] + return "tar", archive, names + + def _extract_archive_member( + self, + kind: str, + archive: Any, + member_name: str, + dest_dir: str + ) -> Optional[str]: + """Stream a single archive entry to dest_dir, preserving its relative path. + + Returns the path of the extracted file, or None if it was skipped. + """ + target = self._safe_member_path(dest_dir, member_name) + if target is None: + self.logger.warning(f"Skipping archive entry with unsafe path: {member_name}") + return None + + parent = os.path.dirname(target) + if parent: + os.makedirs(parent, exist_ok=True) + + if kind == "zip": + source = archive.open(member_name) + else: + source = archive.extractfile(archive.getmember(member_name)) + if source is None: + return None + + try: + with open(target, "wb") as out_file: + shutil.copyfileobj(source, out_file) + finally: + source.close() + + return target + + def process_archive( + self, + service: str, + archive_path: str, + output: Optional[str] = None, + n: int = 10, + generate_ids: bool = False, + consolidate_header: bool = True, + consolidate_citations: bool = False, + include_raw_citations: bool = False, + include_raw_affiliations: bool = False, + tei_coordinates: bool = False, + segment_sentences: bool = False, + force: bool = True, + verbose: bool = False, + flavor: Optional[str] = None, + json_output: bool = False, + markdown_output: bool = False + ) -> None: + """Process the eligible files contained in a zip/tar archive. + + The archive is never fully decompressed: entries are streamed to a + temporary directory in chunks of ``batch_size`` (from the config), each + chunk is sent to GROBID via ``process_batch``, and the temporary files + are removed before the next chunk is extracted. This keeps disk usage + bounded regardless of the archive size. Output files follow the same + flat naming convention as directory processing (one ```` per + result, in ``output``). + """ + start_time = time.time() + self._warn_on_consolidation_timeout(consolidate_citations) + + total_files, processed, errors, skipped = self._process_archive_core( + service, archive_path, output, n, generate_ids, consolidate_header, + consolidate_citations, include_raw_citations, include_raw_affiliations, + tei_coordinates, segment_sentences, force, verbose, flavor, + json_output, markdown_output + ) + + if total_files == 0: + return + + runtime = time.time() - start_time + self._print_processing_summary(processed, errors, skipped, total_files, runtime) + + def _process_archive_core( + self, + service: str, + archive_path: str, + output: Optional[str], + n: int, + generate_ids: bool, + consolidate_header: bool, + consolidate_citations: bool, + include_raw_citations: bool, + include_raw_affiliations: bool, + tei_coordinates: bool, + segment_sentences: bool, + force: bool, + verbose: bool, + flavor: Optional[str], + json_output: bool, + markdown_output: bool + ) -> Tuple[int, int, int, int]: + """Stream and process an archive; return (total, processed, errors, skipped). + + Does not print the final summary (the caller does), so it can be + aggregated with other inputs when resolving a glob pattern. + """ + batch_size_pdf = self.config["batch_size"] + + # Results must survive the temporary extraction directories, so when no + # output is given we default to a directory named after the archive. For + # s3 archives there is no local home, so use the object's basename. + if output is None: + if self._is_s3(archive_path): + output = self._archive_stem(self._s3_basename(archive_path)) + else: + output = self._archive_stem(archive_path) - input_files.append(input_file) + try: + kind, archive, member_names = self._open_archive(archive_path) + except Exception as e: + self.logger.error(f"Could not open archive {archive_path}: {str(e)}") + return 0, 0, 0, 0 + + processed_files_count = 0 + errors_files_count = 0 + skipped_files_count = 0 + total_files = 0 + + try: + eligible_members = [ + name for name in member_names + if self._is_eligible_input(os.path.basename(name), service) + ] + total_files = len(eligible_members) + if total_files == 0: + self.logger.warning(f"No eligible files found in archive {archive_path}") + return 0, 0, 0, 0 + + print(f"Found {total_files} file(s) to process in {archive_path}") + + for chunk_start in range(0, total_files, batch_size_pdf): + chunk = eligible_members[chunk_start:chunk_start + batch_size_pdf] + temp_dir = tempfile.mkdtemp(prefix="grobid_archive_") + try: + extracted_files = [] + for member_name in chunk: + if verbose: + self.logger.info(f"Extracting {member_name} from {archive_path}") + extracted = self._extract_archive_member(kind, archive, member_name, temp_dir) + if extracted is not None: + extracted_files.append(extracted) + + if not extracted_files: + continue + + batch_processed, batch_errors, batch_skipped = self.process_batch( + service, + extracted_files, + temp_dir, + output, + n, + generate_ids, + consolidate_header, + consolidate_citations, + include_raw_citations, + include_raw_affiliations, + tei_coordinates, + segment_sentences, + force, + verbose, + flavor, + json_output, + markdown_output + ) + processed_files_count += batch_processed + errors_files_count += batch_errors + skipped_files_count += batch_skipped + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + finally: + try: + archive.close() + finally: + # ZipFile does not close a file object we passed in (the s3 stream) + stream = getattr(archive, "_grobid_stream", None) + if stream is not None: + stream.close() + + return total_files, processed_files_count, errors_files_count, skipped_files_count + + def _process_remote_files( + self, + service: str, + uris: list, + output: Optional[str], + n: int, + generate_ids: bool, + consolidate_header: bool, + consolidate_citations: bool, + include_raw_citations: bool, + include_raw_affiliations: bool, + tei_coordinates: bool, + segment_sentences: bool, + force: bool, + verbose: bool, + flavor: Optional[str], + json_output: bool, + markdown_output: bool + ) -> Tuple[int, int, int, int]: + """Stream loose remote (s3) files to a temp dir in chunks and process them. + + Returns (total, processed, errors, skipped). Objects are fetched a + batch at a time and deleted before the next chunk, so disk stays bounded. + """ + total = len(uris) + if total == 0: + return 0, 0, 0, 0 + # Remote files have no local home; default output to the current dir. + if output is None: + output = "." + + batch_size_pdf = self.config["batch_size"] + print(f"Found {total} remote file(s) to process") + processed_count = 0 + error_count = 0 + skipped_count = 0 + + for chunk_start in range(0, total, batch_size_pdf): + chunk = uris[chunk_start:chunk_start + batch_size_pdf] + temp_dir = tempfile.mkdtemp(prefix="grobid_s3_") + try: + local_files = [] + for uri in chunk: + if verbose: + self.logger.info(f"Fetching {uri}") + dest = os.path.join(temp_dir, self._s3_basename(uri)) + try: + with self._s3_open(uri) as src, open(dest, "wb") as out_file: + shutil.copyfileobj(src, out_file) + local_files.append(dest) + except Exception as e: + self.logger.error(f"Failed to fetch {uri}: {str(e)}") + error_count += 1 + + if not local_files: + continue - if len(input_files) == batch_size_pdf: batch_processed, batch_errors, batch_skipped = self.process_batch( service, - input_files, - input_path, + local_files, + temp_dir, output, n, generate_ids, @@ -446,49 +1126,13 @@ def process( json_output, markdown_output ) - processed_files_count += batch_processed - errors_files_count += batch_errors - skipped_files_count += batch_skipped - input_files = [] - - # last batch - if len(input_files) > 0: - batch_processed, batch_errors, batch_skipped = self.process_batch( - service, - input_files, - input_path, - output, - n, - generate_ids, - consolidate_header, - consolidate_citations, - include_raw_citations, - include_raw_affiliations, - tei_coordinates, - segment_sentences, - force, - verbose, - flavor, - json_output, - markdown_output - ) - processed_files_count += batch_processed - errors_files_count += batch_errors - skipped_files_count += batch_skipped + processed_count += batch_processed + error_count += batch_errors + skipped_count += batch_skipped + finally: + shutil.rmtree(temp_dir, ignore_errors=True) - runtime = time.time() - start_time - docs_per_second = processed_files_count / runtime if runtime > 0 else 0 - seconds_per_doc = runtime / processed_files_count if processed_files_count > 0 else 0 - - # Log final statistics - always visible - print(f"Processing completed: {processed_files_count} out of {total_files} files processed") - print(f"Errors: {errors_files_count} out of {total_files} files processed") - if skipped_files_count > 0: - print(f"Skipped: {skipped_files_count} out of {total_files} files (already existed, use --force to reprocess)") - - print(f"⏱️ Total runtime: {runtime:.2f} seconds") - print(f"🚀 Speed: {docs_per_second:.2f} documents/second") - print(f" Throughput: {seconds_per_doc:.2f} seconds/document") + return total, processed_count, error_count, skipped_count def process_batch( self, @@ -542,8 +1186,9 @@ def process_batch( json_data = converter.convert_tei_file(filename, stream=False) if json_data: - with open(json_filename_expanded, 'w', encoding='utf8') as json_file: - json.dump(json_data, json_file, indent=2, ensure_ascii=False) + self._write_atomic( + json_filename_expanded, + json.dumps(json_data, indent=2, ensure_ascii=False)) self.logger.debug(f"Successfully created JSON file: {json_filename_expanded}") else: self.logger.warning(f"Failed to convert TEI to JSON for {filename}") @@ -563,8 +1208,7 @@ def process_batch( markdown_data = converter.convert_tei_file(filename) if markdown_data: - with open(markdown_filename_expanded, 'w', encoding='utf8') as markdown_file: - markdown_file.write(markdown_data) + self._write_atomic(markdown_filename_expanded, markdown_data) self.logger.debug(f"Successfully created Markdown file: {markdown_filename_expanded}") else: self.logger.warning(f"Failed to convert TEI to Markdown for {filename}") @@ -606,13 +1250,8 @@ def process_batch( error_count += 1 # writing error file with suffixed error code try: - pathlib.Path(os.path.dirname(filename)).mkdir(parents=True, exist_ok=True) error_filename = filename.replace(".grobid.tei.xml", f"_{status}.txt") - with open(error_filename, 'w', encoding='utf8') as error_file: - if text is not None: - error_file.write(text) - else: - error_file.write("") + self._write_atomic(error_filename, text if text is not None else "") self.logger.info(f"Error details written to {error_filename}") except OSError as e: self.logger.error(f"Failed to write error file {filename}: {str(e)}") @@ -620,9 +1259,7 @@ def process_batch( processed_count += 1 # writing TEI file try: - pathlib.Path(os.path.dirname(filename)).mkdir(parents=True, exist_ok=True) - with open(filename, 'w', encoding='utf8') as tei_file: - tei_file.write(text) + self._write_atomic(filename, text) self.logger.debug(f"Successfully wrote TEI file: {filename}") # Convert to JSON if requested @@ -635,8 +1272,9 @@ def process_batch( json_filename = filename.replace('.grobid.tei.xml', '.json') # Always write JSON file when TEI is written (respects --force behavior) json_filename_expanded = os.path.expanduser(json_filename) - with open(json_filename_expanded, 'w', encoding='utf8') as json_file: - json.dump(json_data, json_file, indent=2, ensure_ascii=False) + self._write_atomic( + json_filename_expanded, + json.dumps(json_data, indent=2, ensure_ascii=False)) self.logger.debug(f"Successfully wrote JSON file: {json_filename_expanded}") else: self.logger.warning(f"Failed to convert TEI to JSON for {filename}") @@ -654,8 +1292,7 @@ def process_batch( markdown_filename = filename.replace('.grobid.tei.xml', '.md') # Always write Markdown file when TEI is written (respects --force behavior) markdown_filename_expanded = os.path.expanduser(markdown_filename) - with open(markdown_filename_expanded, 'w', encoding='utf8') as markdown_file: - markdown_file.write(markdown_data) + self._write_atomic(markdown_filename_expanded, markdown_data) self.logger.debug(f"Successfully wrote Markdown file: {markdown_filename_expanded}") else: self.logger.warning(f"Failed to convert TEI to Markdown for {filename}") @@ -864,7 +1501,12 @@ def main() -> None: parser.add_argument( "--input", default=None, - help="path to the directory containing files to process: PDF or .txt (for processCitationList only, one reference per line), or .xml for patents in ST36" + help="input to process: a directory, a file, a .zip/.tar/.tar.gz archive, a glob pattern (e.g. '**/*.pdf', 'paper*.zip'), or an s3:// object/prefix/glob (requires the 's3' extra). Archives are streamed and never fully decompressed." + ) + parser.add_argument( + "--input-list", + default=None, + help="path to a text file with one input per line (local path, glob or s3:// URI); all are processed together. Lines starting with '#' are ignored." ) parser.add_argument( "--output", @@ -954,6 +1596,7 @@ def main() -> None: args = parser.parse_args() input_path = args.input + input_list = args.input_list config_path = args.config output_path = args.output flavor = args.flavor @@ -1010,12 +1653,31 @@ def main() -> None: logger.error(f"Missing or invalid service '{service}', must be one of {valid_services}") exit(1) + # Build the list of inputs from --input and/or --input-list + inputs = [] + if input_path is not None: + inputs.append(input_path) + if input_list is not None: + try: + with open(input_list, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if line and not line.startswith('#'): + inputs.append(line) + except OSError as e: + logger.error(f"Could not read --input-list {input_list}: {str(e)}") + exit(1) + + if not inputs: + logger.error("No input provided (use --input and/or --input-list)") + exit(1) + start_time = time.time() try: - client.process( + client.process_paths( service, - input_path, + inputs, output=output_path, n=n, generate_ids=generate_ids, diff --git a/pyproject.toml b/pyproject.toml index 7624b63..33ff61b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,11 @@ readme = "Readme.md" dynamic = ['version', "dependencies"] +[project.optional-dependencies] +# Streaming inputs directly from S3 (s3:// URIs/prefixes). smart_open provides +# a seekable, range-streamed reader so remote zips are never fully downloaded. +s3 = ["smart_open[s3]>=6.0", "boto3"] + [tool.setuptools.dynamic] dependencies = {file = ["requirements.txt"]} diff --git a/tests/test_grobid_client.py b/tests/test_grobid_client.py index 7dd6ec3..a46ecdd 100644 --- a/tests/test_grobid_client.py +++ b/tests/test_grobid_client.py @@ -233,41 +233,42 @@ def test_ping_method(self): assert result == (True, 200) - @patch('os.walk') - def test_process_no_files_found(self, mock_walk): + def test_process_no_files_found(self): """Test process method when no eligible files are found.""" - mock_walk.return_value = [('/test/path', [], [])] - - with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): - with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): - client = GrobidClient(check_server=False) - client.logger = Mock() + with tempfile.TemporaryDirectory() as empty_dir: + with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): + with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): + client = GrobidClient(check_server=False) + client.logger = Mock() - client.process('processFulltextDocument', '/test/path') + client.process('processFulltextDocument', empty_dir) - client.logger.warning.assert_called_with('No eligible files found in /test/path') + client.logger.warning.assert_called_with( + f"No eligible files found in input(s): ['{empty_dir}']") - @patch('os.walk') @patch('builtins.print') # Mock print since we use print for statistics - def test_process_with_pdf_files(self, mock_print, mock_walk): - """Test process method with PDF files.""" - mock_walk.return_value = [ - ('/test/path', [], ['doc1.pdf', 'doc2.PDF', 'not_pdf.txt']) - ] - - with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): - with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): - with patch('grobid_client.grobid_client.GrobidClient.process_batch') as mock_batch: - mock_batch.return_value = (2, 0, 0) # Return tuple as expected (processed, errors, skipped) - client = GrobidClient(check_server=False) - client.logger = Mock() + def test_process_with_pdf_files(self, mock_print): + """Test process method with PDF files (directory input).""" + with tempfile.TemporaryDirectory() as input_dir: + for name in ('doc1.pdf', 'doc2.PDF', 'not_pdf.txt'): + with open(os.path.join(input_dir, name), 'wb') as f: + f.write(b'x') + + with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): + with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): + with patch('grobid_client.grobid_client.GrobidClient.process_batch') as mock_batch: + mock_batch.return_value = (2, 0, 0) # (processed, errors, skipped) + client = GrobidClient(check_server=False) + client.logger = Mock() - client.process('processFulltextDocument', '/test/path') + client.process('processFulltextDocument', input_dir) - mock_batch.assert_called_once() - # Check that print was called for statistics - print_calls = [call[0][0] for call in mock_print.call_args_list if 'Found' in call[0][0]] - assert any('Found 2 file(s) to process' in call for call in print_calls) + mock_batch.assert_called_once() + # only the 2 PDFs are batched, the .txt is ignored + batched = mock_batch.call_args.args[1] + assert len(batched) == 2 + print_calls = [call[0][0] for call in mock_print.call_args_list if 'Found' in call[0][0]] + assert any('Found 2 local file(s) to process' in call for call in print_calls) @patch('builtins.open', new_callable=mock_open) @patch('grobid_client.grobid_client.GrobidClient.post') @@ -671,3 +672,350 @@ def test_get_server_url_edge_cases(self, mock_configure_logging, mock_test_serve result = client.get_server_url(service) expected = 'http://localhost:8070/api/processCitationPatentST36' assert result == expected + + +class TestArchiveInput: + """Tests for streaming zip/tar archives as input (process_archive).""" + + def _client(self, batch_size=2): + with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): + with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): + client = GrobidClient(check_server=False) + client.logger = Mock() + client.config['batch_size'] = batch_size + return client + + @staticmethod + def _make_zip(path, entries): + import zipfile + with zipfile.ZipFile(path, 'w') as z: + for name, data in entries.items(): + z.writestr(name, data) + + @staticmethod + def _make_targz(path, entries, work): + import tarfile + with tarfile.open(path, 'w:gz') as t: + for name, data in entries.items(): + member_path = os.path.join(work, os.path.basename(name)) + with open(member_path, 'wb') as f: + f.write(data) + t.add(member_path, arcname=name) + + def _run(self, client, archive, output): + """Run archive processing with a fake GROBID post; return set of temp dirs used.""" + temp_dirs = set() + + def fake_post(url, files=None, data=None, headers=None, timeout=None): + temp_dirs.add(os.path.dirname(files['input'][0])) + resp = Mock() + resp.text = 'ok' + return (resp, 200) + + with patch.object(GrobidClient, 'post', side_effect=fake_post): + client.process('processFulltextDocument', archive, output=output, force=True) + return temp_dirs + + @staticmethod + def _tei_outputs(output_dir): + found = [] + for root, _, files in os.walk(output_dir): + for f in files: + if f.endswith('.grobid.tei.xml'): + found.append(f) + return sorted(found) + + def test_is_archive(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'x.zip') + self._make_zip(zip_path, {'a.pdf': b'%PDF'}) + assert client._is_archive(zip_path) is True + assert client._is_archive(d) is False # directory + assert client._is_archive(os.path.join(d, 'missing.zip')) is False + + def test_archive_stem(self): + client = self._client() + assert client._archive_stem('/x/docs.tar.gz') == '/x/docs' + assert client._archive_stem('/x/docs.tgz') == '/x/docs' + assert client._archive_stem('/x/docs.zip') == '/x/docs' + + def test_safe_member_path_blocks_traversal(self): + client = self._client() + dest = os.path.join('some', 'dest') + # traversal and absolute paths are neutralized to stay under dest + assert client._safe_member_path(dest, '../../etc/passwd') == os.path.join(dest, 'etc', 'passwd') + assert client._safe_member_path(dest, '/abs/evil.pdf') == os.path.join(dest, 'abs', 'evil.pdf') + assert client._safe_member_path(dest, '') is None + assert client._safe_member_path(dest, '.') is None + + def test_process_zip_streams_all_pdfs(self): + client = self._client(batch_size=2) + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'docs.zip') + self._make_zip(zip_path, { + 'a.pdf': b'%PDF-a', + 'sub/b.pdf': b'%PDF-b', + 'c.PDF': b'%PDF-c', + 'ignore.txt': b'not a pdf', + }) + out = os.path.join(d, 'out') + temp_dirs = self._run(client, zip_path, out) + + # all 3 PDFs processed, the .txt ignored + assert self._tei_outputs(out) == ['a.grobid.tei.xml', 'b.grobid.tei.xml', 'c.grobid.tei.xml'] + # 3 files with batch_size 2 => 2 chunks => distinct temp dirs, all cleaned up + assert len(temp_dirs) >= 2 + assert all(not os.path.exists(td) for td in temp_dirs) + + def test_process_targz(self): + client = self._client(batch_size=10) + with tempfile.TemporaryDirectory() as d: + tar_path = os.path.join(d, 'docs.tar.gz') + self._make_targz(tar_path, {'x.pdf': b'%PDF-x', 'nested/y.pdf': b'%PDF-y'}, d) + out = os.path.join(d, 'out') + temp_dirs = self._run(client, tar_path, out) + assert self._tei_outputs(out) == ['x.grobid.tei.xml', 'y.grobid.tei.xml'] + assert all(not os.path.exists(td) for td in temp_dirs) + + def test_process_routes_archive_to_core(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'docs.zip') + self._make_zip(zip_path, {'a.pdf': b'%PDF'}) + with patch.object(GrobidClient, '_process_archive_core', return_value=(1, 1, 0, 0)) as mock_core: + client.process('processFulltextDocument', zip_path, output=os.path.join(d, 'o')) + mock_core.assert_called_once() + assert mock_core.call_args.args[1] == zip_path + + def test_process_zip_default_output_named_after_archive(self): + client = self._client(batch_size=10) + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'mydocs.zip') + self._make_zip(zip_path, {'a.pdf': b'%PDF'}) + self._run(client, zip_path, None) # no output -> defaults to + assert self._tei_outputs(os.path.join(d, 'mydocs')) == ['a.grobid.tei.xml'] + + def test_empty_archive_warns(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'empty.zip') + self._make_zip(zip_path, {'notes.txt': b'no pdfs here'}) + with patch.object(GrobidClient, 'process_batch') as mock_batch: + client.process('processFulltextDocument', zip_path, output=os.path.join(d, 'o')) + mock_batch.assert_not_called() + client.logger.warning.assert_called() + + +class TestGlobInput: + """Tests for glob-pattern input resolution (--input as a glob).""" + + def _client(self, batch_size=50): + with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): + with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): + client = GrobidClient(check_server=False) + client.logger = Mock() + client.config['batch_size'] = batch_size + return client + + @staticmethod + def _zip(path, entries): + import zipfile + with zipfile.ZipFile(path, 'w') as z: + for name, data in entries.items(): + z.writestr(name, data) + + @staticmethod + def _tei_outputs(output_dir): + found = [] + for root, _, files in os.walk(output_dir): + for f in files: + if f.endswith('.grobid.tei.xml'): + found.append(f) + return sorted(found) + + def _run(self, client, pattern, output): + def fake_post(url, files=None, data=None, headers=None, timeout=None): + resp = Mock() + resp.text = 'ok' + return (resp, 200) + with patch.object(GrobidClient, 'post', side_effect=fake_post): + client.process('processFulltextDocument', pattern, output=output, force=True) + + def test_resolve_input_paths_plain_and_glob(self): + client = self._client() + # plain path (no magic) returned as-is even if missing + assert client._resolve_input_paths('/nope/x.zip') == ['/nope/x.zip'] + with tempfile.TemporaryDirectory() as d: + for n in ('paper1.zip', 'paper2.zip', 'other.zip'): + open(os.path.join(d, n), 'wb').close() + matches = client._resolve_input_paths(os.path.join(d, 'paper*.zip')) + assert [os.path.basename(m) for m in matches] == ['paper1.zip', 'paper2.zip'] + + def test_glob_matches_multiple_archives(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + self._zip(os.path.join(d, 'paper1.zip'), {'a.pdf': b'%PDF-a'}) + self._zip(os.path.join(d, 'paper2.zip'), {'b.pdf': b'%PDF-b'}) + self._zip(os.path.join(d, 'skip.zip'), {'c.pdf': b'%PDF-c'}) + out = os.path.join(d, 'out') + self._run(client, os.path.join(d, 'paper*.zip'), out) + # only paper1/paper2 archives, skip.zip excluded by the pattern + assert self._tei_outputs(out) == ['a.grobid.tei.xml', 'b.grobid.tei.xml'] + + def test_glob_recursive_pdfs_across_subdirs(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + os.makedirs(os.path.join(d, 'sub1')) + os.makedirs(os.path.join(d, 'sub2')) + open(os.path.join(d, 'sub1', 'a.pdf'), 'wb').close() + open(os.path.join(d, 'sub2', 'b.pdf'), 'wb').close() + open(os.path.join(d, 'sub2', 'note.txt'), 'wb').close() + out = os.path.join(d, 'out') + self._run(client, os.path.join(d, '**', '*.pdf'), out) + assert self._tei_outputs(out) == ['a.grobid.tei.xml', 'b.grobid.tei.xml'] + + def test_glob_no_match_warns(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + client.process('processFulltextDocument', os.path.join(d, 'nothing*.zip'), + output=os.path.join(d, 'o')) + client.logger.warning.assert_called() + assert "No files match" in client.logger.warning.call_args[0][0] + + def test_common_base_is_ancestor(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + f1 = os.path.join(d, 'x', 'a.pdf') + f2 = os.path.join(d, 'y', 'b.pdf') + os.makedirs(os.path.dirname(f1)); os.makedirs(os.path.dirname(f2)) + open(f1, 'wb').close(); open(f2, 'wb').close() + base = client._common_base([f1, f2]) + assert os.path.isdir(base) + assert f1.startswith(base) and f2.startswith(base) + + +class TestAtomicWrite: + """A killed task must never leave a partial output that resume accepts. + + process_batch decides a document is done with os.path.isfile() alone, so a + TEI truncated by an OOM kill is skipped forever on subsequent runs. These + tests pin the two properties that prevent it: a completed write is whole, + and a failed write leaves nothing behind at all. + """ + + def _client(self): + with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): + with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): + client = GrobidClient(check_server=False) + client.logger = Mock() + return client + + @staticmethod + def _leftovers(directory): + """Temp files the writer may have leaked (they are dot-prefixed).""" + return [n for n in os.listdir(directory) if n.startswith('.') and n.endswith('.tmp')] + + def test_write_lands_content(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + dest = os.path.join(d, 'a.grobid.tei.xml') + client._write_atomic(dest, 'content') + with open(dest, encoding='utf8') as fh: + assert fh.read() == 'content' + assert self._leftovers(d) == [] + + def test_write_creates_missing_parents(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + dest = os.path.join(d, 'sub', 'dir', 'a.grobid.tei.xml') + client._write_atomic(dest, 'x') + assert os.path.isfile(dest) + + def test_write_overwrites_existing(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + dest = os.path.join(d, 'a.grobid.tei.xml') + client._write_atomic(dest, 'old and much longer') + client._write_atomic(dest, 'new') + with open(dest, encoding='utf8') as fh: + assert fh.read() == 'new' + + def test_failed_write_leaves_no_destination(self): + """The whole point: a write that dies part-way must not create the output. + + Simulates the real failure -- some bytes reach the disk, then the + process dies -- by writing a prefix and raising, as an OOM kill would. + """ + client = self._client() + real_fdopen = os.fdopen + + class PartialWriter: + def __init__(self, handle): + self._handle = handle + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + self._handle.close() + return False + + def write(self, text): + self._handle.write(text[:4]) # bytes hit the disk... + raise RuntimeError('killed mid-write') # ...then the task dies + + with tempfile.TemporaryDirectory() as d: + dest = os.path.join(d, 'a.grobid.tei.xml') + with patch('os.fdopen', lambda fd, *a, **kw: PartialWriter(real_fdopen(fd, *a, **kw))): + with pytest.raises(RuntimeError): + client._write_atomic(dest, 'content') + + assert not os.path.exists(dest), \ + "a partial write created the destination -- resume would skip it forever" + assert self._leftovers(d) == [], "temp file was left behind" + + def test_failed_write_preserves_previous_content(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + dest = os.path.join(d, 'a.grobid.tei.xml') + client._write_atomic(dest, 'good') + with patch('os.replace', side_effect=OSError('boom')): + with pytest.raises(OSError): + client._write_atomic(dest, 'bad') + with open(dest, encoding='utf8') as fh: + assert fh.read() == 'good' + assert self._leftovers(d) == [] + + def test_temp_name_is_not_counted_as_output(self): + """grobid_stream.sh counts *.grobid.tei.xml and *_[0-9]*.txt via find. + + An in-flight temp file must match neither, or the tallies move while a + write is happening. + """ + import fnmatch + client = self._client() + seen = {} + real_mkstemp = tempfile.mkstemp + + def spy(*args, **kwargs): + fd, path = real_mkstemp(*args, **kwargs) + seen['name'] = os.path.basename(path) + return fd, path + + with tempfile.TemporaryDirectory() as d: + with patch('tempfile.mkstemp', side_effect=spy): + client._write_atomic(os.path.join(d, 'a.grobid.tei.xml'), 'x') + assert not fnmatch.fnmatch(seen['name'], '*.grobid.tei.xml') + assert not fnmatch.fnmatch(seen['name'], '*_[0-9]*.txt') + + def test_write_uses_normal_permissions(self): + """mkstemp creates 0600; TEIs on shared scratch must stay group-readable.""" + client = self._client() + with tempfile.TemporaryDirectory() as d: + reference = os.path.join(d, 'reference.txt') + with open(reference, 'w') as fh: # what the old code produced + fh.write('x') + dest = os.path.join(d, 'a.grobid.tei.xml') + client._write_atomic(dest, 'x') + assert (os.stat(dest).st_mode & 0o777) == (os.stat(reference).st_mode & 0o777) diff --git a/tests/test_s3.py b/tests/test_s3.py new file mode 100644 index 0000000..1ff1512 --- /dev/null +++ b/tests/test_s3.py @@ -0,0 +1,148 @@ +""" +Tests for streaming inputs from S3 (the optional 's3' extra). + +These use moto to mock S3 and mock GrobidClient.post so no GROBID server or real +AWS is needed. +""" +import io +import os +import zipfile + +import pytest +from unittest.mock import Mock, patch + +boto3 = pytest.importorskip("boto3") +pytest.importorskip("smart_open") +pytest.importorskip("moto") +try: + from moto import mock_aws +except ImportError: # moto < 5 + from moto import mock_s3 as mock_aws + +from grobid_client.grobid_client import GrobidClient + +BUCKET = "test-bucket" +REGION = "us-east-1" + + +@pytest.fixture(autouse=True) +def _aws_env(monkeypatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "testing") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "testing") + monkeypatch.setenv("AWS_SECURITY_TOKEN", "testing") + monkeypatch.setenv("AWS_SESSION_TOKEN", "testing") + monkeypatch.setenv("AWS_DEFAULT_REGION", REGION) + + +def _client(batch_size=2): + with patch.object(GrobidClient, "_test_server_connection"): + with patch.object(GrobidClient, "_configure_logging"): + c = GrobidClient(check_server=False) + c.logger = Mock() + c.config["batch_size"] = batch_size + return c + + +def _fake_post(url, files=None, data=None, headers=None, timeout=None): + resp = Mock() + resp.text = "ok" + return (resp, 200) + + +def _zip_bytes(entries): + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as z: + for name, data in entries.items(): + z.writestr(name, data) + return buf.getvalue() + + +def _tei(out): + return sorted(f for _, _, fs in os.walk(out) for f in fs if f.endswith(".grobid.tei.xml")) + + +def test_split_and_basename(): + c = _client() + assert c._split_s3("s3://bucket/a/b/c.zip") == ("bucket", "a/b/c.zip") + assert c._s3_basename("s3://bucket/a/b/c.zip") == "c.zip" + assert c._is_s3("s3://bucket/x") is True + assert c._is_s3("/local/x") is False + + +def test_resolve_single_object_needs_no_listing(): + # a concrete object key is returned as-is (no S3 call at all) + c = _client() + assert c._resolve_s3_paths("s3://bucket/a/b/file.zip") == ["s3://bucket/a/b/file.zip"] + + +@mock_aws +def test_resolve_prefix_and_glob(): + s3 = boto3.client("s3", region_name=REGION) + s3.create_bucket(Bucket=BUCKET) + for k in ["p/0000.zip", "p/0001.zip", "p/readme.txt", "q/x.zip"]: + s3.put_object(Bucket=BUCKET, Key=k, Body=b"x") + + c = _client() + assert c._resolve_s3_paths(f"s3://{BUCKET}/p/") == [ + f"s3://{BUCKET}/p/0000.zip", f"s3://{BUCKET}/p/0001.zip", f"s3://{BUCKET}/p/readme.txt", + ] + assert c._resolve_s3_paths(f"s3://{BUCKET}/p/*.zip") == [ + f"s3://{BUCKET}/p/0000.zip", f"s3://{BUCKET}/p/0001.zip", + ] + + +@mock_aws +def test_process_s3_zip_range_streamed(tmp_path): + s3 = boto3.client("s3", region_name=REGION) + s3.create_bucket(Bucket=BUCKET) + s3.put_object(Bucket=BUCKET, Key="arch/docs.zip", Body=_zip_bytes({ + "0000001.pdf": b"%PDF-a", + "sub/0000002.pdf": b"%PDF-b", + "note.txt": b"not a pdf", + })) + c = _client() + out = str(tmp_path / "out") + with patch.object(GrobidClient, "post", side_effect=_fake_post): + c.process("processFulltextDocument", f"s3://{BUCKET}/arch/docs.zip", output=out, force=True) + # both PDFs processed, .txt ignored + assert _tei(out) == ["0000001.grobid.tei.xml", "0000002.grobid.tei.xml"] + + +@mock_aws +def test_process_s3_loose_pdfs_glob(tmp_path): + s3 = boto3.client("s3", region_name=REGION) + s3.create_bucket(Bucket=BUCKET) + for k in ["pdfs/0000001.pdf", "pdfs/0000002.pdf", "pdfs/skip.txt"]: + s3.put_object(Bucket=BUCKET, Key=k, Body=b"%PDF") + c = _client() + out = str(tmp_path / "out") + with patch.object(GrobidClient, "post", side_effect=_fake_post): + c.process("processFulltextDocument", f"s3://{BUCKET}/pdfs/*.pdf", output=out, force=True) + assert _tei(out) == ["0000001.grobid.tei.xml", "0000002.grobid.tei.xml"] + + +@mock_aws +def test_process_paths_mixed_local_and_s3(tmp_path): + s3 = boto3.client("s3", region_name=REGION) + s3.create_bucket(Bucket=BUCKET) + s3.put_object(Bucket=BUCKET, Key="a.zip", Body=_zip_bytes({"0000009.pdf": b"%PDF-z"})) + local_pdf = tmp_path / "local.pdf" + local_pdf.write_bytes(b"%PDF-l") + + c = _client() + out = str(tmp_path / "out") + with patch.object(GrobidClient, "post", side_effect=_fake_post): + c.process_paths( + "processFulltextDocument", + [str(local_pdf), f"s3://{BUCKET}/a.zip"], + output=out, force=True, + ) + assert _tei(out) == ["0000009.grobid.tei.xml", "local.grobid.tei.xml"] + + +def test_missing_extra_raises_helpful_error(): + """If smart_open isn't importable, a clear install hint is raised.""" + c = _client() + with patch.dict("sys.modules", {"smart_open": None}): + with pytest.raises(ImportError, match=r"pip install grobid-client-python\[s3\]"): + c._s3_open("s3://bucket/key.zip")