Support dict as the input of delete_batch (#30241)

lostluck · shunping · web-flow · commit 8978c8079c20 · 2024-02-06T16:30:20.000-08:00
This smooths the code path of feeding output of list_prefix()
to delete_batch().

Co-authored-by: Shunping Huang &lt;shunping@google.com&gt;
diff --git a/sdks/python/apache_beam/io/gcp/gcsio.py b/sdks/python/apache_beam/io/gcp/gcsio.py
@@ -208,7 +208,8 @@ def delete_batch(self, paths):
     """Deletes the objects at the given GCS paths.
 
     Args:
-      paths: List of GCS file path patterns in the form gs://<bucket>/<name>,
+      paths: List of GCS file path patterns or Dict with GCS file path patterns
+             as keys. The patterns are in the form gs://<bucket>/<name>, but
              not to exceed MAX_BATCH_OPERATION_SIZE in length.
 
     Returns: List of tuples of (path, exception) in the same order as the
@@ -217,6 +218,7 @@ def delete_batch(self, paths):
     """
     final_results = []
     s = 0
+    if not isinstance(paths, list): paths = list(iter(paths))
     while s < len(paths):
       if (s + MAX_BATCH_OPERATION_SIZE) < len(paths):
         current_paths = paths[s:s + MAX_BATCH_OPERATION_SIZE]
