diff --git a/Readme.md b/Readme.md index 31f6f19..2082196 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,26 @@ grobid_client --server https://grobid.example.com --input ~/citations.txt proces # Force reprocessing with sentence segmentation and JSON output grobid_client --input ~/docs --force --segment_sentences --json processFulltextDocument + +# Process PDFs directly from a zip or tar.gz archive (streamed, not fully decompressed) +grobid_client --input ~/papers.zip --output ~/results processFulltextDocument +grobid_client --input ~/papers.tar.gz --output ~/results processFulltextDocument + +# --input also accepts glob patterns (quote them so the shell does not expand them) +grobid_client --input "~/papers/*.zip" --output ~/results processFulltextDocument # many archives +grobid_client --input "~/data/**/*.pdf" --output ~/results processFulltextDocument # PDFs in subdirectories ``` +> [!NOTE] +> `--input` accepts a directory, a single file, an **archive**, or a **glob pattern**: +> - **Archives** (`.zip`, `.tar`, `.tar.gz`/`.tgz`, `.tar.bz2`/`.tbz2`) are streamed: eligible entries are extracted in +> chunks of `batch_size` to a temporary directory, sent to GROBID, written to `--output`, and deleted before the next +> chunk. The archive is never fully decompressed, so disk usage stays bounded. If `--output` is omitted, results go to a +> directory named after the archive (e.g. `papers.zip` → `papers/`). +> - **Glob patterns** (`paper.zip`, `paper*.zip`, `**/paper*.zip`, `**/*.pdf`, …) are expanded with `**` recursion; each +> match is handled by type (archive → streamed, directory → recursed, file → processed). Quote the pattern so your shell +> passes it through to the client unexpanded. + ### Python Library #### Basic Usage diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index a525b9f..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 @@ -24,6 +25,10 @@ import requests import pathlib import logging +import shutil +import tarfile +import tempfile +import zipfile from typing import Tuple import copy @@ -45,6 +50,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', @@ -369,7 +379,6 @@ def process( markdown_output=False ): 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 @@ -378,108 +387,416 @@ 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 filename.endswith(".pdf") or filename.endswith(".PDF") or \ - (service == 'processCitationList' and ( - filename.endswith(".txt") or filename.endswith(".TXT"))) or \ - (service == 'processCitationPatentST36' and ( - filename.endswith(".xml") or filename.endswith(".XML"))): - full_path = os.sep.join([dirpath, filename]) - all_input_files.append(full_path) - - # Log total files found - total_files = len(all_input_files) - if total_files == 0: - self.logger.warning(f"No eligible files found in {input_path}") + 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) + 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 - print(f"Found {total_files} file(s) to process") - input_files = [] + 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"⏱️ 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 + return processed_files_count, errors_files_count, skipped_files_count + + def _is_eligible_input(self, filename, service): + """Return True if a file name is a valid input for the given service.""" + if filename.endswith(".pdf") or filename.endswith(".PDF"): + return True + if service == 'processCitationList' and ( + filename.endswith(".txt") or filename.endswith(".TXT")): + return True + if service == 'processCitationPatentST36' and ( + filename.endswith(".xml") or filename.endswith(".XML")): + return True + return False + + def _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() + 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 - 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") + 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: + 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 0, 0, 0, 0 + + processed_files_count = 0 + errors_files_count = 0 + skipped_files_count = 0 + total_files = 0 + + try: + eligible_members = [ + name for name in member_names + if self._is_eligible_input(os.path.basename(name), service) + ] + total_files = len(eligible_members) + if total_files == 0: + self.logger.warning(f"No eligible files found in archive {archive_path}") + return 0, 0, 0, 0 + + print(f"Found {total_files} file(s) to process in {archive_path}") + + for chunk_start in range(0, total_files, batch_size_pdf): + chunk = eligible_members[chunk_start:chunk_start + batch_size_pdf] + temp_dir = tempfile.mkdtemp(prefix="grobid_archive_") + try: + extracted_files = [] + for member_name in chunk: + if verbose: + self.logger.info(f"Extracting {member_name} from {archive_path}") + extracted = self._extract_archive_member(kind, archive, member_name, temp_dir) + if extracted is not None: + extracted_files.append(extracted) + + if not extracted_files: + continue + + batch_processed, batch_errors, batch_skipped = self.process_batch( + service, + extracted_files, + temp_dir, + output, + n, + generate_ids, + consolidate_header, + consolidate_citations, + include_raw_citations, + include_raw_affiliations, + tei_coordinates, + segment_sentences, + force, + verbose, + flavor, + json_output, + markdown_output + ) + processed_files_count += batch_processed + errors_files_count += batch_errors + skipped_files_count += batch_skipped + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + finally: + archive.close() + + return total_files, processed_files_count, errors_files_count, skipped_files_count def process_batch( self, @@ -855,7 +1172,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..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') @@ -671,3 +672,224 @@ def test_get_server_url_edge_cases(self, mock_configure_logging, mock_test_serve result = client.get_server_url(service) expected = 'http://localhost:8070/api/processCitationPatentST36' assert result == expected + + +class TestArchiveInput: + """Tests for streaming zip/tar archives as input (process_archive).""" + + def _client(self, batch_size=2): + with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): + with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): + client = GrobidClient(check_server=False) + client.logger = Mock() + client.config['batch_size'] = batch_size + return client + + @staticmethod + def _make_zip(path, entries): + import zipfile + with zipfile.ZipFile(path, 'w') as z: + for name, data in entries.items(): + z.writestr(name, data) + + @staticmethod + def _make_targz(path, entries, work): + import tarfile + with tarfile.open(path, 'w:gz') as t: + for name, data in entries.items(): + member_path = os.path.join(work, os.path.basename(name)) + with open(member_path, 'wb') as f: + f.write(data) + t.add(member_path, arcname=name) + + def _run(self, client, archive, output): + """Run archive processing with a fake GROBID post; return set of temp dirs used.""" + temp_dirs = set() + + def fake_post(url, files=None, data=None, headers=None, timeout=None): + temp_dirs.add(os.path.dirname(files['input'][0])) + resp = Mock() + resp.text = 'ok' + return (resp, 200) + + with patch.object(GrobidClient, 'post', side_effect=fake_post): + client.process('processFulltextDocument', archive, output=output, force=True) + return temp_dirs + + @staticmethod + def _tei_outputs(output_dir): + found = [] + for root, _, files in os.walk(output_dir): + for f in files: + if f.endswith('.grobid.tei.xml'): + found.append(f) + return sorted(found) + + def test_is_archive(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'x.zip') + self._make_zip(zip_path, {'a.pdf': b'%PDF'}) + assert client._is_archive(zip_path) is True + assert client._is_archive(d) is False # directory + assert client._is_archive(os.path.join(d, 'missing.zip')) is False + + def test_archive_stem(self): + client = self._client() + assert client._archive_stem('/x/docs.tar.gz') == '/x/docs' + assert client._archive_stem('/x/docs.tgz') == '/x/docs' + assert client._archive_stem('/x/docs.zip') == '/x/docs' + + def test_safe_member_path_blocks_traversal(self): + client = self._client() + dest = os.path.join('some', 'dest') + # traversal and absolute paths are neutralized to stay under dest + assert client._safe_member_path(dest, '../../etc/passwd') == os.path.join(dest, 'etc', 'passwd') + assert client._safe_member_path(dest, '/abs/evil.pdf') == os.path.join(dest, 'abs', 'evil.pdf') + assert client._safe_member_path(dest, '') is None + assert client._safe_member_path(dest, '.') is None + + def test_process_zip_streams_all_pdfs(self): + client = self._client(batch_size=2) + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'docs.zip') + self._make_zip(zip_path, { + 'a.pdf': b'%PDF-a', + 'sub/b.pdf': b'%PDF-b', + 'c.PDF': b'%PDF-c', + 'ignore.txt': b'not a pdf', + }) + out = os.path.join(d, 'out') + temp_dirs = self._run(client, zip_path, out) + + # all 3 PDFs processed, the .txt ignored + assert self._tei_outputs(out) == ['a.grobid.tei.xml', 'b.grobid.tei.xml', 'c.grobid.tei.xml'] + # 3 files with batch_size 2 => 2 chunks => distinct temp dirs, all cleaned up + assert len(temp_dirs) >= 2 + assert all(not os.path.exists(td) for td in temp_dirs) + + def test_process_targz(self): + client = self._client(batch_size=10) + with tempfile.TemporaryDirectory() as d: + tar_path = os.path.join(d, 'docs.tar.gz') + self._make_targz(tar_path, {'x.pdf': b'%PDF-x', 'nested/y.pdf': b'%PDF-y'}, d) + out = os.path.join(d, 'out') + temp_dirs = self._run(client, tar_path, out) + assert self._tei_outputs(out) == ['x.grobid.tei.xml', 'y.grobid.tei.xml'] + assert all(not os.path.exists(td) for td in temp_dirs) + + def test_process_routes_archive_to_core(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'docs.zip') + self._make_zip(zip_path, {'a.pdf': b'%PDF'}) + with patch.object(GrobidClient, '_process_archive_core', return_value=(1, 1, 0, 0)) as mock_core: + client.process('processFulltextDocument', zip_path, output=os.path.join(d, 'o')) + mock_core.assert_called_once() + assert mock_core.call_args.args[1] == zip_path + + def test_process_zip_default_output_named_after_archive(self): + client = self._client(batch_size=10) + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'mydocs.zip') + self._make_zip(zip_path, {'a.pdf': b'%PDF'}) + self._run(client, zip_path, None) # no output -> defaults to + assert self._tei_outputs(os.path.join(d, 'mydocs')) == ['a.grobid.tei.xml'] + + def test_empty_archive_warns(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'empty.zip') + self._make_zip(zip_path, {'notes.txt': b'no pdfs here'}) + with patch.object(GrobidClient, 'process_batch') as mock_batch: + client.process('processFulltextDocument', zip_path, output=os.path.join(d, 'o')) + mock_batch.assert_not_called() + client.logger.warning.assert_called() + + +class TestGlobInput: + """Tests for glob-pattern input resolution (--input as a glob).""" + + def _client(self, batch_size=50): + with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): + with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): + client = GrobidClient(check_server=False) + client.logger = Mock() + client.config['batch_size'] = batch_size + return client + + @staticmethod + def _zip(path, entries): + import zipfile + with zipfile.ZipFile(path, 'w') as z: + for name, data in entries.items(): + z.writestr(name, data) + + @staticmethod + def _tei_outputs(output_dir): + found = [] + for root, _, files in os.walk(output_dir): + for f in files: + if f.endswith('.grobid.tei.xml'): + found.append(f) + return sorted(found) + + def _run(self, client, pattern, output): + def fake_post(url, files=None, data=None, headers=None, timeout=None): + resp = Mock() + resp.text = 'ok' + return (resp, 200) + with patch.object(GrobidClient, 'post', side_effect=fake_post): + client.process('processFulltextDocument', pattern, output=output, force=True) + + def test_resolve_input_paths_plain_and_glob(self): + client = self._client() + # plain path (no magic) returned as-is even if missing + assert client._resolve_input_paths('/nope/x.zip') == ['/nope/x.zip'] + with tempfile.TemporaryDirectory() as d: + for n in ('paper1.zip', 'paper2.zip', 'other.zip'): + open(os.path.join(d, n), 'wb').close() + matches = client._resolve_input_paths(os.path.join(d, 'paper*.zip')) + assert [os.path.basename(m) for m in matches] == ['paper1.zip', 'paper2.zip'] + + def test_glob_matches_multiple_archives(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + self._zip(os.path.join(d, 'paper1.zip'), {'a.pdf': b'%PDF-a'}) + self._zip(os.path.join(d, 'paper2.zip'), {'b.pdf': b'%PDF-b'}) + self._zip(os.path.join(d, 'skip.zip'), {'c.pdf': b'%PDF-c'}) + out = os.path.join(d, 'out') + self._run(client, os.path.join(d, 'paper*.zip'), out) + # only paper1/paper2 archives, skip.zip excluded by the pattern + assert self._tei_outputs(out) == ['a.grobid.tei.xml', 'b.grobid.tei.xml'] + + def test_glob_recursive_pdfs_across_subdirs(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + os.makedirs(os.path.join(d, 'sub1')) + os.makedirs(os.path.join(d, 'sub2')) + open(os.path.join(d, 'sub1', 'a.pdf'), 'wb').close() + open(os.path.join(d, 'sub2', 'b.pdf'), 'wb').close() + open(os.path.join(d, 'sub2', 'note.txt'), 'wb').close() + out = os.path.join(d, 'out') + self._run(client, os.path.join(d, '**', '*.pdf'), out) + assert self._tei_outputs(out) == ['a.grobid.tei.xml', 'b.grobid.tei.xml'] + + def test_glob_no_match_warns(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + client.process('processFulltextDocument', os.path.join(d, 'nothing*.zip'), + output=os.path.join(d, 'o')) + client.logger.warning.assert_called() + assert "No files match" in client.logger.warning.call_args[0][0] + + def test_common_base_is_ancestor(self): + client = self._client() + with tempfile.TemporaryDirectory() as d: + f1 = os.path.join(d, 'x', 'a.pdf') + f2 = os.path.join(d, 'y', 'b.pdf') + os.makedirs(os.path.dirname(f1)); os.makedirs(os.path.dirname(f2)) + open(f1, 'wb').close(); open(f2, 'wb').close() + base = client._common_base([f1, f2]) + assert os.path.isdir(base) + assert f1.startswith(base) and f2.startswith(base)