Skip to content

Commit eb2702f

Browse files
duckduckgrayduckduckduckgrayduck
committed
Handle force_ocr, re-factor upload_directory, and record tests
1 parent 02a0b60 commit eb2702f

80 files changed

Lines changed: 3835 additions & 3060 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.github/workflows/main.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ jobs:
1313
runs-on: ubuntu-latest
1414
strategy:
1515
matrix:
16-
python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
16+
python-version: ["3.10", "3.11"]
1717
steps:
1818
- name: Check out code
1919
uses: actions/checkout@v2

‎Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ test-create-short: test-clean-short
2020

2121
# run all tests using the pre-recorded HTTP requests - will fail on missing request
2222
test:
23-
pytest --record-mode=none --block-network
23+
pytest --record-mode=none --block-network -vv
2424

2525
# run tests using pre-recorded HTTP requets if they exist, and recording them if missing
2626
test-dev:

‎docs/conf.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,16 +48,16 @@
4848

4949
# General information about the project.
5050
project = "documentcloud"
51-
copyright = "2023, MuckRock Foundation"
51+
copyright = "2025, MuckRock Foundation"
5252

5353
# The version info for the project you're documenting, acts as replacement for
5454
# |version| and |release|, also used in various other places throughout the
5555
# built documents.
5656
#
5757
# The short X.Y version.
58-
version = "4.3"
58+
version = "4.5"
5959
# The full version, including alpha/beta/rc tags.
60-
release = "4.3.0"
60+
release = "4.5.0"
6161

6262
# The language for content autogenerated by Sphinx. Refer to documentation
6363
# for a list of supported languages.

‎documentcloud/documents.py‎

Lines changed: 111 additions & 134 deletions
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,11 @@ def __str__(self):
7474
def __getattr__(self, attr):
7575
"""Generate methods for fetching resources"""
7676
p_image = re.compile(
77-
r"^get_(?P<size>thumbnail|small|normal|large|xlarge)_image_url(?P<list>_list)?$"
77+
r"^get_"
78+
r"(?P<size>thumbnail|small|normal|large|xlarge)_image_url"
79+
r"(?P<list>_list)?$"
7880
)
81+
7982
get = attr.startswith("get_")
8083
url = attr.endswith("_url")
8184
text = attr.endswith("_text")
@@ -230,9 +233,15 @@ def get_errors(self):
230233

231234
return all_results
232235

233-
def process(self):
234-
"""Reprocess the document"""
235-
self._client.post(f"{self.api_path}/{self.id}/process/")
236+
def process(self, **kwargs):
237+
"""Process the document, used on upload and for reprocessing"""
238+
payload = {}
239+
if "force_ocr" in kwargs:
240+
payload["force_ocr"] = kwargs["force_ocr"]
241+
if "ocr_engine" in kwargs:
242+
payload["ocr_engine"] = kwargs["ocr_engine"]
243+
244+
self._client.post(f"{self.api_path}/{self.id}/process/", json=payload)
236245

237246

238247
class DocumentClient(BaseAPIClient):
@@ -310,6 +319,7 @@ def _format_upload_parameters(self, name, **kwargs):
310319
"title",
311320
"data",
312321
"force_ocr",
322+
"ocr_engine",
313323
"projects",
314324
"delayed_index",
315325
"revision_control",
@@ -333,21 +343,55 @@ def _format_upload_parameters(self, name, **kwargs):
333343

334344
return params
335345

346+
def _extract_ocr_options(self, kwargs):
347+
"""
348+
Extract and validate OCR options from kwargs.
349+
350+
Returns:
351+
force_ocr (bool)
352+
ocr_engine (str)
353+
"""
354+
force_ocr = kwargs.pop("force_ocr", False)
355+
ocr_engine = kwargs.pop("ocr_engine", "tess4")
356+
357+
if not isinstance(force_ocr, bool):
358+
raise ValueError("force_ocr must be a boolean")
359+
360+
if ocr_engine and ocr_engine not in ("tess4", "textract"):
361+
raise ValueError(
362+
"ocr_engine must be either 'tess4' for tesseract or 'textract'"
363+
)
364+
365+
return force_ocr, ocr_engine
366+
336367
def _get_title(self, name):
337368
"""Get the default title for a document from its path"""
338369
return name.split(os.sep)[-1].rsplit(".", 1)[0]
339370

340371
def _upload_url(self, file_url, **kwargs):
341372
"""Upload a document from a publicly accessible URL"""
373+
# extract process-related args
374+
force_ocr, ocr_engine = self._extract_ocr_options(kwargs)
375+
376+
# create the document
342377
params = self._format_upload_parameters(file_url, **kwargs)
343378
params["file_url"] = file_url
379+
if force_ocr:
380+
params["force_ocr"] = force_ocr
381+
params["ocr_engine"] = ocr_engine
344382
response = self.client.post("documents/", json=params)
345-
return Document(self.client, response.json())
383+
create_json = response.json()
384+
385+
# wrap in Document object
386+
doc = Document(self.client, create_json)
387+
388+
return doc
346389

347390
def _upload_file(self, file_, **kwargs):
348391
"""Upload a document directly"""
349392
# create the document
350-
force_ocr = kwargs.pop("force_ocr", False)
393+
force_ocr, ocr_engine = self._extract_ocr_options(kwargs)
394+
351395
params = self._format_upload_parameters(file_.name, **kwargs)
352396
response = self.client.post("documents/", json=params)
353397

@@ -357,12 +401,12 @@ def _upload_file(self, file_, **kwargs):
357401
response = requests_retry_session().put(presigned_url, data=file_.read())
358402

359403
# begin processing the document
360-
doc_id = create_json["id"]
361-
response = self.client.post(
362-
f"documents/{doc_id}/process/", json={"force_ocr": force_ocr}
363-
)
404+
doc = Document(self.client, create_json)
364405

365-
return Document(self.client, create_json)
406+
# begin processing
407+
doc.process(force_ocr=force_ocr, ocr_engine=ocr_engine)
408+
409+
return doc
366410

367411
def _collect_files(self, path, extensions):
368412
"""Find the paths to files with specified extensions under a directory"""
@@ -379,165 +423,98 @@ def _collect_files(self, path, extensions):
379423

380424
def upload_directory(self, path, handle_errors=False, extensions=".pdf", **kwargs):
381425
"""Upload files with specified extensions in a directory"""
382-
# pylint: disable=too-many-locals, too-many-branches
383-
384-
# Do not set the same title for all documents
426+
# pylint:disable=too-many-locals
385427
kwargs.pop("title", None)
386428

387-
# If extensions are specified as None, it will check for all supported
388-
# filetypes.
389429
if extensions is None:
390430
extensions = SUPPORTED_EXTENSIONS
391-
392-
# Convert single extension to a list if provided
393431
if extensions and not isinstance(extensions, list):
394432
extensions = [extensions]
395-
396-
# Checks to see if the extensions are supported, raises an error if not.
397433
invalid_extensions = set(extensions) - set(SUPPORTED_EXTENSIONS)
398434
if invalid_extensions:
399435
raise ValueError(
400436
f"Invalid extensions provided: {', '.join(invalid_extensions)}"
401437
)
402438

403-
# Loop through the path and get all the files with matching extensions
404439
path_list = self._collect_files(path, extensions)
405-
406440
logger.info(
407441
"Upload directory on %s: Found %d files to upload", path, len(path_list)
408442
)
409443

410-
# Upload all the files using the bulk API to reduce the number
411-
# of API calls and improve performance
412444
obj_list = []
445+
force_ocr, ocr_engine = self._extract_ocr_options(kwargs)
413446
params = self._format_upload_parameters("", **kwargs)
447+
414448
for i, file_paths in enumerate(grouper(path_list, BULK_LIMIT)):
415-
# Grouper will put None's on the end of the last group
416449
file_paths = [p for p in file_paths if p is not None]
417-
418450
logger.info("Uploading group %d:\n%s", i + 1, "\n".join(file_paths))
419451

420-
# Create the documents
421-
logger.info("Creating the documents...")
422-
try:
423-
response = self.client.post(
424-
"documents/",
425-
json=[
426-
merge_dicts(
427-
params,
428-
{
429-
"title": self._get_title(p),
430-
"original_extension": os.path.splitext(
431-
os.path.basename(p)
432-
)[1]
433-
.lower()
434-
.lstrip("."),
435-
},
436-
)
437-
for p in file_paths
438-
],
439-
)
440-
except (APIError, RequestException) as exc:
441-
if handle_errors:
442-
logger.info(
443-
"Error creating the following documents: %s\n%s",
444-
exc,
445-
"\n".join(file_paths),
446-
)
447-
continue
448-
else:
449-
raise
452+
create_json = self._create_documents(file_paths, params, handle_errors)
453+
sorted_create_json = sorted(create_json, key=lambda j: j["title"])
454+
sorted_file_paths = sorted(file_paths, key=self._get_title)
455+
obj_list.extend(sorted_create_json)
456+
presigned_urls = [j["presigned_url"] for j in sorted_create_json]
450457

451-
# Upload the files directly to storage
452-
create_json = response.json()
453-
obj_list.extend(create_json)
454-
presigned_urls = [j["presigned_url"] for j in create_json]
455-
for url, file_path in zip(presigned_urls, file_paths):
456-
logger.info("Uploading %s to S3...", file_path)
457-
try:
458-
with open(file_path, "rb") as file:
459-
response = requests_retry_session().put(url, data=file.read())
460-
self.client.raise_for_status(response)
461-
except (APIError, RequestException) as exc:
462-
if handle_errors:
463-
logger.info(
464-
"Error uploading the following document: %s %s",
465-
exc,
466-
file_path,
467-
)
468-
continue
469-
else:
470-
raise
471-
472-
# Begin processing the documents
473-
logger.info("Processing the documents...")
474-
doc_ids = [j["id"] for j in create_json]
475-
try:
476-
response = self.client.post("documents/process/", json={"ids": doc_ids})
477-
except (APIError, RequestException) as exc:
478-
if handle_errors:
479-
logger.info(
480-
"Error creating the following documents: %s\n%s",
481-
exc,
482-
"\n".join(file_paths),
483-
)
484-
continue
485-
else:
486-
raise
458+
self._upload_files_to_s3(sorted_file_paths, presigned_urls, handle_errors)
459+
self._process_documents(create_json, force_ocr, ocr_engine, handle_errors)
487460

488461
logger.info("Upload directory complete")
489-
490-
# Pass back the list of documents
491462
return [Document(self.client, d) for d in obj_list]
492463

493-
def upload_urls(self, url_list, handle_errors=False, **kwargs):
494-
"""Upload documents from a list of URLs"""
495-
496-
# Do not set the same title for all documents
497-
kwargs.pop("title", None)
498-
499-
obj_list = []
500-
params = self._format_upload_parameters("", **kwargs)
501-
for i, url_group in enumerate(grouper(url_list, BULK_LIMIT)):
502-
# Grouper will put None's on the end of the last group
503-
url_group = [url for url in url_group if url is not None]
504-
505-
logger.info("Uploading group %d: %s", i + 1, "\n".join(url_group))
506-
507-
# Create the documents
508-
logger.info("Creating the documents...")
509-
try:
510-
response = self.client.post(
511-
"documents/",
512-
json=[
513-
merge_dicts(
514-
params,
515-
{
516-
"title": self._get_title(url),
517-
"file_url": url,
518-
},
519-
)
520-
for url in url_group
521-
],
464+
def _create_documents(self, file_paths, params, handle_errors):
465+
body = [
466+
merge_dicts(
467+
params,
468+
{
469+
"title": self._get_title(p),
470+
"original_extension": os.path.splitext(os.path.basename(p))[1]
471+
.lower()
472+
.lstrip("."),
473+
},
474+
)
475+
for p in sorted(file_paths)
476+
]
477+
try:
478+
response = self.client.post("documents/", json=body)
479+
except (APIError, RequestException) as exc:
480+
if handle_errors:
481+
logger.info(
482+
"Error creating the following documents: %s\n%s",
483+
exc,
484+
"\n".join(file_paths),
522485
)
486+
return []
487+
else:
488+
raise
489+
return response.json()
490+
491+
def _upload_files_to_s3(self, file_paths, presigned_urls, handle_errors):
492+
for url, file_path in zip(presigned_urls, file_paths):
493+
logger.info("Uploading %s to S3...", file_path)
494+
try:
495+
with open(file_path, "rb") as f:
496+
response = requests_retry_session().put(url, data=f.read())
497+
self.client.raise_for_status(response)
523498
except (APIError, RequestException) as exc:
524499
if handle_errors:
525500
logger.info(
526-
"Error creating the following documents: %s\n%s",
527-
str(exc),
528-
"\n".join(url_group),
501+
"Error uploading the following document: %s %s", exc, file_path
529502
)
530-
continue
531503
else:
532504
raise
533505

534-
create_json = response.json()
535-
obj_list.extend(create_json)
536-
537-
logger.info("Upload URLs complete")
538-
539-
# Pass back the list of documents
540-
return [Document(self.client, d) for d in obj_list]
506+
def _process_documents(self, create_json, force_ocr, ocr_engine, handle_errors):
507+
payload = [
508+
{"id": j["id"], "force_ocr": force_ocr, "ocr_engine": ocr_engine}
509+
for j in create_json
510+
]
511+
try:
512+
self.client.post("documents/process/", json=payload)
513+
except (APIError, RequestException) as exc:
514+
if handle_errors:
515+
logger.info("Error processing documents: %s", exc)
516+
else:
517+
raise
541518

542519

543520
class Mention:

‎setup.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
setup(
99
name="python-documentcloud",
10-
version="4.4.1",
10+
version="4.5.0",
1111
description="A simple Python wrapper for the DocumentCloud API",
1212
author="Mitchell Kotler",
1313
author_email="mitch@muckrock.com",

0 commit comments

Comments
 (0)