feat: Added Document.list_from_gcs() method to get multiple wrapped documents from a GCS directory.

Fixes #214

- Note: `from_gcs()` takes in a GCS directory, but it only works for a single sharded document from a single input document source.
- In a GA release, it would be a better practice to have `from_gcs()` take in any GCS directory and output a list of Wrapped Documents. But this would be a backwards-incompatible change now.
  - Not sure if it's possible/advisable to have two possible return types for `from_gcs()` and just have it return a list when there are multiple Wrapped Documents?

Author: holtskinner

google/cloud/documentai_toolbox/wrappers/document.py

@@ -481,6 +481,35 @@ def from_gcs(
             gcs_input_uri=gcs_input_uri,
         )
 
+    @classmethod
+    def list_from_gcs(
+        cls: Type["Document"],
+        gcs_bucket_name: str,
+        gcs_prefix: str,
+    ) -> List["Document"]:
+        r"""Loads a list of Documents from Cloud Storage.
+
+        Args:
+            gcs_bucket_name (str):
+                Required. The gcs bucket.
+
+                Format: Given `gs://{bucket_name}/{optional_folder}/{operation_id}/` where `gcs_bucket_name={bucket_name}`.
+            gcs_prefix (str):
+                Required. The prefix to the location of the target folder.
+
+                Format: Given `gs://{bucket_name}/{optional_folder}/{target_folder}` where `gcs_prefix={optional_folder}/{target_folder}`.
+        Returns:
+            List[Document]:
+                A List of documents from gcs.
+        """
+        return [
+            Document.from_gcs(gcs_bucket_name=gcs_bucket_name, gcs_prefix=directory)
+            for directory, files in gcs_utilities.list_gcs_document_tree(
+                gcs_bucket_name, gcs_prefix
+            ).items()
+            if files != [""]
+        ]
+
     @classmethod
     def from_batch_process_metadata(
         cls: Type["Document"], metadata: documentai.BatchProcessMetadata

tests/unit/test_document.py

@@ -30,6 +30,7 @@
 import pytest
 
 from google.cloud import documentai
+from google.cloud.storage import Blob, Bucket
 from google.cloud.documentai_toolbox import document, gcs_utilities
 
 
@@ -397,6 +398,37 @@ def test_document_from_gcs_with_unordered_shards(get_bytes_unordered_files_mock)
         assert page.page_number == page_index + 1
 
 
+@mock.patch("google.cloud.documentai_toolbox.utilities.gcs_utilities.storage")
+def test_document_list_from_gcs_with_multiple_input_files(
+    mock_storage,
+    get_bytes_multiple_directories_mock,
+):
+    client = mock_storage.Client.return_value
+
+    mock_bucket = mock.Mock()
+
+    client.Bucket.return_value = mock_bucket
+
+    client.list_blobs.return_value = [
+        Blob(name="documentai/output/123456789/1/test_shard1.json", bucket=None),
+        Blob(name="documentai/output/123456789/1/test_shard2.json", bucket=None),
+        Blob(name="documentai/output/123456789/2/test_shard3.json", bucket=None),
+    ]
+    documents = document.Document.list_from_gcs(
+        gcs_bucket_name="test-directory", gcs_prefix="documentai/output/123456789/"
+    )
+    get_bytes_multiple_directories_mock.assert_called()
+    assert get_bytes_multiple_directories_mock.call_count == 2
+
+    assert len(documents) == 2
+
+    assert documents[0].gcs_bucket_name == "test-directory"
+    assert documents[0].gcs_prefix == "documentai/output/123456789/1"
+
+    assert documents[1].gcs_bucket_name == "test-directory"
+    assert documents[1].gcs_prefix == "documentai/output/123456789/2"
+
+
 def test_document_from_batch_process_metadata_with_multiple_input_files(
     get_bytes_multiple_directories_mock,
 ):