From 88040f2508e6783a3d6fcdaad2203d9651f01a3d Mon Sep 17 00:00:00 2001 From: Luca Foppiano Date: Thu, 23 Jul 2026 19:36:22 +0200 Subject: [PATCH 1/3] Stream PDFs from zip/tar archives without full decompression Add process_archive(), used automatically when --input points to a .zip/.tar/.tar.gz (.tgz/.tar.bz2/.tbz2) archive. Eligible entries are read out of the archive in chunks of batch_size, each chunk is extracted to a temporary directory, sent to GROBID via the existing process_batch (so concurrency, TEI/JSON/Markdown output, --force/skip and error handling are reused), and the temporary files are removed before the next chunk is extracted. The archive is never fully decompressed, so disk usage stays bounded regardless of its size. - zip via zipfile, tar/tar.gz/tar.bz2 via tarfile ('r:*') - entries are streamed member-by-member (zipfile.open / tarfile.extractfile) - archive paths are sanitized to prevent path traversal (zip slip) - when --output is omitted, results go to a directory named after the archive - directory-input eligibility check refactored into _is_eligible_input and shared Documented in the Readme and covered by unit tests (zip, tar.gz, chunking, cleanup, default output, traversal guard, delegation). --- Readme.md | 12 ++ grobid_client/grobid_client.py | 242 ++++++++++++++++++++++++++++++++- tests/test_grobid_client.py | 133 ++++++++++++++++++ 3 files changed, 381 insertions(+), 6 deletions(-) diff --git a/Readme.md b/Readme.md index 31f6f19..bcb4bb6 100644 --- a/Readme.md +++ b/Readme.md @@ -33,6 +33,7 @@ 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 ## 📋 Prerequisites @@ -166,8 +167,19 @@ 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 ``` +> [!NOTE] +> When `--input` points to a `.zip`, `.tar`, `.tar.gz`/`.tgz` (or `.tar.bz2`/`.tbz2`) archive, the client streams the +> eligible entries out of it in chunks of `batch_size` (from the config): each chunk is extracted to a temporary +> directory, sent to GROBID, written to `--output`, and deleted before the next chunk is extracted. The archive is never +> fully decompressed, so disk usage stays bounded regardless of its size. If `--output` is omitted, results are written to +> a directory named after the archive (e.g. `papers.zip` → `papers/`). + ### Python Library #### Basic Usage diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index a525b9f..a376614 100644 --- a/grobid_client/grobid_client.py +++ b/grobid_client/grobid_client.py @@ -24,6 +24,10 @@ import requests import pathlib import logging +import shutil +import tarfile +import tempfile +import zipfile from typing import Tuple import copy @@ -45,6 +49,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,6 +377,29 @@ def process( json_output=False, markdown_output=False ): + # If the input is a zip/tar archive, stream its entries out one chunk at + # a time instead of walking a directory. This never fully decompresses + # the archive and keeps disk usage bounded. + if input_path is not None and self._is_archive(input_path): + return self.process_archive( + 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, + ) + start_time = time.time() batch_size_pdf = self.config["batch_size"] @@ -382,11 +414,7 @@ def process( 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"))): + if self._is_eligible_input(filename, service): full_path = os.sep.join([dirpath, filename]) all_input_files.append(full_path) @@ -481,6 +509,208 @@ def process( print(f"🚀 Speed: {docs_per_second:.2f} documents/second") print(f" Throughput: {seconds_per_doc:.2f} seconds/document") + 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 _is_archive(self, path): + """Return True if path is an existing zip/tar archive file.""" + if not os.path.isfile(path): + return False + lower = path.lower() + return any(lower.endswith(ext) for ext in self.ARCHIVE_EXTENSIONS) + + 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). + """ + 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() + batch_size_pdf = self.config["batch_size"] + self._warn_on_consolidation_timeout(consolidate_citations) + + # Results must survive the temporary extraction directories, so when no + # output is given we default to a directory named after the archive. + if output is None: + output = self._archive_stem(archive_path) + + try: + kind, archive, member_names = self._open_archive(archive_path) + except (zipfile.BadZipFile, tarfile.TarError, OSError) as e: + self.logger.error(f"Could not open archive {archive_path}: {str(e)}") + return + + 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 + + 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: + archive.close() + + if total_files == 0: + return + + 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 + + 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") + def process_batch( self, service, @@ -855,7 +1085,7 @@ 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="path to the directory - or a .zip/.tar/.tar.gz archive - containing files to process: PDF or .txt (for processCitationList only, one reference per line), or .xml for patents in ST36. Archives are streamed and never fully decompressed." ) parser.add_argument( "--output", diff --git a/tests/test_grobid_client.py b/tests/test_grobid_client.py index 7dd6ec3..e73440e 100644 --- a/tests/test_grobid_client.py +++ b/tests/test_grobid_client.py @@ -671,3 +671,136 @@ 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_delegates_archive_to_process_archive(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') as mock_archive: + client.process('processFulltextDocument', zip_path, output=os.path.join(d, 'o')) + mock_archive.assert_called_once() + assert mock_archive.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() From 74b004631948f2e405e739c22776db7b2a91b315 Mon Sep 17 00:00:00 2001 From: Luca Foppiano Date: Fri, 24 Jul 2026 13:06:41 +0200 Subject: [PATCH 2/3] Support glob patterns for --input and use pathlib for directory traversal --input now accepts shell-style glob patterns (with recursive **), e.g. 'paper.zip' (one file), 'paper*.zip' (many), '**/paper*.zip' (subdirectories) or '**/*.pdf'. Each match is dispatched by type: archives are streamed, directories are recursed, eligible files are processed directly, and the results of all matches are aggregated into a single summary. - resolve --input via glob (has_magic + recursive=True, ~ expansion); a plain path is returned unchanged for backward compatibility - directory traversal refactored from os.walk to pathlib.Path.rglob - factor the batching loop, stats summary and archive streaming into reusable helpers (_run_file_batches, _print_processing_summary, _process_archive_core) so directory, loose-file and archive inputs share one code path - loose files matched by a glob are batched together under their common base Documented in the Readme and covered by unit tests (multi-archive glob, recursive **/*.pdf, mixed matches, no-match warning, common-base helper). --- Readme.md | 17 +- grobid_client/grobid_client.py | 313 +++++++++++++++++++++------------ tests/test_grobid_client.py | 153 ++++++++++++---- 3 files changed, 333 insertions(+), 150 deletions(-) diff --git a/Readme.md b/Readme.md index bcb4bb6..2082196 100644 --- a/Readme.md +++ b/Readme.md @@ -171,14 +171,21 @@ grobid_client --input ~/docs --force --segment_sentences --json processFulltextD # 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] -> When `--input` points to a `.zip`, `.tar`, `.tar.gz`/`.tgz` (or `.tar.bz2`/`.tbz2`) archive, the client streams the -> eligible entries out of it in chunks of `batch_size` (from the config): each chunk is extracted to a temporary -> directory, sent to GROBID, written to `--output`, and deleted before the next chunk is extracted. The archive is never -> fully decompressed, so disk usage stays bounded regardless of its size. If `--output` is omitted, results are written to -> a directory named after the archive (e.g. `papers.zip` → `papers/`). +> `--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. ### Python Library diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index a376614..f138650 100644 --- a/grobid_client/grobid_client.py +++ b/grobid_client/grobid_client.py @@ -17,6 +17,7 @@ import os import json import argparse +import glob import time import concurrent.futures import ntpath @@ -377,31 +378,7 @@ def process( json_output=False, markdown_output=False ): - # If the input is a zip/tar archive, stream its entries out one chunk at - # a time instead of walking a directory. This never fully decompresses - # the archive and keeps disk usage bounded. - if input_path is not None and self._is_archive(input_path): - return self.process_archive( - 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, - ) - 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 @@ -410,104 +387,190 @@ def process( # 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 self._is_eligible_input(filename, service): - full_path = os.sep.join([dirpath, filename]) - all_input_files.append(full_path) + if input_path is None: + self.logger.warning("No input path provided") + return - # 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}") + # input_path may be a plain directory/file/archive or a glob pattern + # (e.g. "paper*.zip", "**/*.pdf"). Resolve it to concrete paths. + matched_paths = self._resolve_input_paths(input_path) + if not matched_paths: + self.logger.warning(f"No files match input '{input_path}'") + return + + # Partition matches into archives (streamed) and plain filesystem files + # (directories are expanded to their eligible files). + archive_paths = [] + fs_files = [] + for path in matched_paths: + if 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: + self.logger.warning(f"No eligible files found in input '{input_path}'") return - # Counters for processing statistics (initialize before early return) processed_files_count = 0 errors_files_count = 0 skipped_files_count = 0 + total_files = 0 + + # Plain files gathered from directories and/or loose glob matches + if fs_files: + print(f"Found {len(fs_files)} file(s) to process") + batch_processed, batch_errors, batch_skipped = 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 += batch_processed + errors_files_count += batch_errors + skipped_files_count += batch_skipped + total_files += len(fs_files) + + # Archives are streamed entry-by-entry, one batch-sized chunk at a time + for archive_path in archive_paths: + arc_total, arc_processed, arc_errors, arc_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 + ) + processed_files_count += arc_processed + errors_files_count += arc_errors + skipped_files_count += arc_skipped + total_files += arc_total + + if total_files == 0: + self.logger.warning(f"No eligible files found in input '{input_path}'") + 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 path into a sorted list of concrete paths. + + Supports shell-style glob patterns (including the recursive ``**``) and + ``~`` expansion. A plain path without glob metacharacters is returned + as-is (so callers can still handle a missing path themselves). + """ + 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) + + 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"Found {total_files} file(s) to process") - input_files = [] + 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. - for input_file in all_input_files: - # Extract just the filename for verbose logging - filename = os.path.basename(input_file) + 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(input_files) == batch_size_pdf: + if len(batch) == batch_size_pdf: 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 + 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 - input_files = [] + batch = [] - # last batch - if len(input_files) > 0: + if batch: 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 + 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 - 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 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.""" @@ -622,9 +685,47 @@ def process_archive( result, in ``output``). """ start_time = time.time() - batch_size_pdf = self.config["batch_size"] 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. if output is None: @@ -634,7 +735,7 @@ def process_archive( kind, archive, member_names = self._open_archive(archive_path) except (zipfile.BadZipFile, tarfile.TarError, OSError) as e: self.logger.error(f"Could not open archive {archive_path}: {str(e)}") - return + return 0, 0, 0, 0 processed_files_count = 0 errors_files_count = 0 @@ -649,7 +750,7 @@ def process_archive( total_files = len(eligible_members) if total_files == 0: self.logger.warning(f"No eligible files found in archive {archive_path}") - return + return 0, 0, 0, 0 print(f"Found {total_files} file(s) to process in {archive_path}") @@ -695,21 +796,7 @@ def process_archive( finally: archive.close() - if total_files == 0: - return - - 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 - - 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_files, processed_files_count, errors_files_count, skipped_files_count def process_batch( self, diff --git a/tests/test_grobid_client.py b/tests/test_grobid_client.py index e73440e..95d7d92 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 '{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 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') @@ -777,15 +778,15 @@ def test_process_targz(self): 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_delegates_archive_to_process_archive(self): + 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') as mock_archive: + 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_archive.assert_called_once() - assert mock_archive.call_args.args[1] == zip_path + 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) @@ -804,3 +805,91 @@ def test_empty_archive_warns(self): 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) From 5400cf2676fb4fb4e81e5ffd1f757ca6a1121d92 Mon Sep 17 00:00:00 2001 From: Luca Foppiano Date: Sun, 26 Jul 2026 09:18:25 +0200 Subject: [PATCH 3/3] Stream inputs from S3 (s3://) via optional [s3] extra Add s3:// support to --input (and a new --input-list manifest). An s3 zip is range-streamed with smart_open (only the central directory and the requested entries are fetched - the object is never fully downloaded); loose remote PDFs are fetched a batch at a time to a temp dir. Mixed manifests (local + glob + s3) are supported and aggregated into one summary. - refactor process() -> process_paths(list of inputs); process() delegates - _resolve_input_paths handles s3 object / prefix / glob (list_objects_v2 + fnmatch) - _open_archive range-streams s3 zips (smart_open seekable stream); the stream is closed after use; s3 tar is rejected (not range-streamable) - new _process_remote_files streams loose s3 objects in bounded chunks - --input-list reads a file of paths (local/glob/s3, '#' comments) - s3 deps (smart_open[s3], boto3) are an optional 'pip install ...[s3]' extra, lazily imported with a clear install hint if missing - credentials use the standard AWS chain (env / ~/.aws / IAM) Documented in the Readme; covered by moto-backed tests (single object, prefix, glob, zip range-streaming, loose PDFs, mixed manifest, missing-extra error). --- Readme.md | 15 ++ grobid_client/grobid_client.py | 358 ++++++++++++++++++++++++++++----- pyproject.toml | 5 + tests/test_grobid_client.py | 4 +- tests/test_s3.py | 148 ++++++++++++++ 5 files changed, 483 insertions(+), 47 deletions(-) create mode 100644 tests/test_s3.py diff --git a/Readme.md b/Readme.md index 2082196..4832e50 100644 --- a/Readme.md +++ b/Readme.md @@ -34,6 +34,7 @@ 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 - **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 +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 @@ -186,6 +190,17 @@ grobid_client --input "~/data/**/*.pdf" --output ~/results processFulltextDocu > - **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 diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index f138650..dab1714 100644 --- a/grobid_client/grobid_client.py +++ b/grobid_client/grobid_client.py @@ -17,6 +17,7 @@ import os import json import argparse +import fnmatch import glob import time import concurrent.futures @@ -378,32 +379,71 @@ 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() - # 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) - if input_path is None: - self.logger.warning("No input path provided") - return - - # input_path may be a plain directory/file/archive or a glob pattern - # (e.g. "paper*.zip", "**/*.pdf"). Resolve it to concrete paths. - matched_paths = self._resolve_input_paths(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 '{input_path}'") + self.logger.warning(f"No files match input(s): {inputs}") return - # Partition matches into archives (streamed) and plain filesystem files - # (directories are expanded to their eligible files). + # 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_archive(path): + 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)) @@ -412,8 +452,8 @@ def process( else: self.logger.debug(f"Skipping input (not an eligible file/dir/archive): {path}") - if not fs_files and not archive_paths: - self.logger.warning(f"No eligible files found in input '{input_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 processed_files_count = 0 @@ -421,35 +461,48 @@ def process( skipped_files_count = 0 total_files = 0 - # Plain files gathered from directories and/or loose glob matches + # Local files gathered from directories and/or loose glob matches if fs_files: - print(f"Found {len(fs_files)} file(s) to process") - batch_processed, batch_errors, batch_skipped = self._run_file_batches( + 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 += batch_processed - errors_files_count += batch_errors - skipped_files_count += batch_skipped + processed_files_count += bp + errors_files_count += be + skipped_files_count += bs total_files += len(fs_files) - # Archives are streamed entry-by-entry, one batch-sized chunk at a time + # 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: - arc_total, arc_processed, arc_errors, arc_skipped = self._process_archive_core( + 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 += arc_processed - errors_files_count += arc_errors - skipped_files_count += arc_skipped - total_files += arc_total + 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 '{input_path}'") + self.logger.warning(f"No eligible files found in input(s): {inputs}") return runtime = time.time() - start_time @@ -458,12 +511,15 @@ def process( ) def _resolve_input_paths(self, input_path): - """Resolve an input path into a sorted list of concrete paths. + """Resolve an input into a sorted list of concrete paths. - Supports shell-style glob patterns (including the recursive ``**``) and - ``~`` expansion. A plain path without glob metacharacters is returned - as-is (so callers can still handle a missing path themselves). + 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)) @@ -493,6 +549,84 @@ def _common_base(self, files): 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 + + def _s3_basename(self, uri): + """Return the last path component of an s3:// key.""" + return self._split_s3(uri)[1].rsplit("/", 1)[-1] + + 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 @@ -584,13 +718,15 @@ def _is_eligible_input(self, filename, service): return True return False - def _is_archive(self, path): - """Return True if path is an existing zip/tar archive file.""" - if not os.path.isfile(path): - 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() @@ -616,7 +752,20 @@ 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("/")] @@ -727,13 +876,17 @@ def _process_archive_core( 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. + # 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: - output = self._archive_stem(archive_path) + 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 (zipfile.BadZipFile, tarfile.TarError, OSError) as e: + except Exception as e: self.logger.error(f"Could not open archive {archive_path}: {str(e)}") return 0, 0, 0, 0 @@ -794,10 +947,100 @@ def _process_archive_core( finally: shutil.rmtree(temp_dir, ignore_errors=True) finally: - archive.close() + 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 + + batch_processed, batch_errors, batch_skipped = self.process_batch( + service, + local_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_count += batch_processed + error_count += batch_errors + skipped_count += batch_skipped + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + return total, processed_count, error_count, skipped_count + def process_batch( self, service, @@ -1172,7 +1415,12 @@ def main(): parser.add_argument( "--input", default=None, - help="path to the directory - or a .zip/.tar/.tar.gz archive - containing files to process: PDF or .txt (for processCitationList only, one reference per line), or .xml for patents in ST36. Archives are streamed and never fully decompressed." + 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", @@ -1262,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 @@ -1318,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 95d7d92..5df1312 100644 --- a/tests/test_grobid_client.py +++ b/tests/test_grobid_client.py @@ -244,7 +244,7 @@ def test_process_no_files_found(self): client.process('processFulltextDocument', empty_dir) client.logger.warning.assert_called_with( - f"No eligible files found in input '{empty_dir}'") + f"No eligible files found in input(s): ['{empty_dir}']") @patch('builtins.print') # Mock print since we use print for statistics def test_process_with_pdf_files(self, mock_print): @@ -268,7 +268,7 @@ def test_process_with_pdf_files(self, mock_print): 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 file(s) to process' in call for call in print_calls) + 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') 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")