diff --git a/Readme.md b/Readme.md index 31f6f19..4832e50 100644 --- a/Readme.md +++ b/Readme.md @@ -33,6 +33,8 @@ concurrent processing capabilities for PDF documents, reference strings, and pat - **Sentence Segmentation**: Layout-aware sentence segmentation capabilities - **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 +- **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 @@ -57,6 +59,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 @@ -166,8 +171,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 a525b9f..dab1714 100644 --- a/grobid_client/grobid_client.py +++ b/grobid_client/grobid_client.py @@ -17,6 +17,8 @@ import os import json import argparse +import fnmatch +import glob import time import concurrent.futures import ntpath @@ -24,6 +26,10 @@ import requests import pathlib import logging +import shutil +import tarfile +import tempfile +import zipfile from typing import Tuple import copy @@ -45,6 +51,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 = { 'grobid_server': 'http://localhost:8070', @@ -368,60 +379,645 @@ def process( json_output=False, markdown_output=False ): + 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, + inputs, + output=None, + n=10, + generate_ids=False, + consolidate_header=True, + consolidate_citations=False, + include_raw_citations=False, + include_raw_affiliations=False, + tei_coordinates=False, + segment_sentences=False, + force=True, + verbose=False, + flavor=None, + json_output=False, + markdown_output=False + ): + """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): + """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, service): + """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): + """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): + """Return True if path is an s3:// URI.""" + return isinstance(path, str) and path.startswith("s3://") + + @staticmethod + def _split_s3(uri): + """Split an s3://bucket/key URI into (bucket, key).""" + bucket, _, key = uri[len("s3://"):].partition("/") + return bucket, key - print(f"Found {total_files} file(s) to process") - input_files = [] + def _s3_basename(self, uri): + """Return the last path component of an s3:// key.""" + return self._split_s3(uri)[1].rsplit("/", 1)[-1] - for input_file in all_input_files: - # Extract just the filename for verbose logging - filename = os.path.basename(input_file) + def _import_smart_open(self): + 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): + 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): + """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): + """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, errors, skipped, total, runtime): + """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, + 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 + ): + """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") - input_files.append(input_file) + 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, service): + """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): + """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): + """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): + """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, arcname): + """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): + """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. + """ + 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, archive, member_name, dest_dir): + """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, + archive_path, + output=None, + n=10, + generate_ids=False, + consolidate_header=True, + consolidate_citations=False, + include_raw_citations=False, + include_raw_affiliations=False, + tei_coordinates=False, + segment_sentences=False, + force=True, + verbose=False, + flavor=None, + json_output=False, + markdown_output=False + ): + """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, + 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 + ): + """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) + + 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, + uris, + 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 + ): + """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, @@ -437,49 +1033,13 @@ def process( json_output, markdown_output ) - processed_files_count += batch_processed - errors_files_count += batch_errors - skipped_files_count += batch_skipped - input_files = [] + processed_count += batch_processed + error_count += batch_errors + skipped_count += batch_skipped + finally: + shutil.rmtree(temp_dir, ignore_errors=True) - # 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 - - 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, @@ -855,7 +1415,12 @@ def main(): 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", @@ -945,6 +1510,7 @@ def main(): args = parser.parse_args() input_path = args.input + input_list = args.input_list config_path = args.config output_path = args.output flavor = args.flavor @@ -1001,12 +1567,31 @@ def main(): 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 9dadbd9..913680b 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..5df1312 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,224 @@ 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) 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")