From 88040f2508e6783a3d6fcdaad2203d9651f01a3d Mon Sep 17 00:00:00 2001 From: Luca Foppiano Date: Thu, 23 Jul 2026 19:36:22 +0200 Subject: [PATCH 1/2] 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/2] 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)