samples: Add snippets for upload/download from file-like object/stream (googleapis#709)

andrewsg · web-flow · commit 126d94c63843 · 2022-02-02T16:30:57.000-08:00
* samples: Add snippets for upload/download from file-like object/stream

* change sample region tags to be uniform with other libraries

* lint
diff --git a/samples/snippets/snippets_test.py b/samples/snippets/snippets_test.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import io
 import os
 import tempfile
 import time
@@ -44,6 +45,7 @@
 import storage_download_file
 import storage_download_into_memory
 import storage_download_public_file
+import storage_download_to_stream
 import storage_enable_bucket_lifecycle_management
 import storage_enable_versioning
 import storage_generate_signed_post_policy_v4
@@ -68,6 +70,7 @@
 import storage_set_metadata
 import storage_upload_file
 import storage_upload_from_memory
+import storage_upload_from_stream
 import storage_upload_with_kms_key
 
 KMS_KEY = os.environ["CLOUD_KMS_KEY"]
@@ -204,6 +207,17 @@ def test_upload_blob_from_memory(test_bucket, capsys):
     assert "Hello, is it me you're looking for?" in out
 
 
+def test_upload_blob_from_stream(test_bucket, capsys):
+    file_obj = io.StringIO()
+    file_obj.write("This is test data.")
+    storage_upload_from_stream.upload_blob_from_stream(
+        test_bucket.name, file_obj, "test_upload_blob"
+    )
+    out, _ = capsys.readouterr()
+
+    assert "Stream data uploaded to {}".format("test_upload_blob") in out
+
+
 def test_upload_blob_with_kms(test_bucket):
     with tempfile.NamedTemporaryFile() as source_file:
         source_file.write(b"test")
@@ -247,6 +261,20 @@ def test_download_blob_into_memory(test_blob, capsys):
     assert "Hello, is it me you're looking for?" in out
 
 
+def test_download_blob_to_stream(test_blob, capsys):
+    file_obj = io.BytesIO()
+    storage_download_to_stream.download_blob_to_stream(
+        test_blob.bucket.name, test_blob.name, file_obj
+    )
+    out, _ = capsys.readouterr()
+
+    file_obj.seek(0)
+    content = file_obj.read()
+
+    assert "Downloaded blob" in out
+    assert b"Hello, is it me you're looking for?" in content
+
+
 def test_blob_metadata(test_blob, capsys):
     storage_get_metadata.blob_metadata(test_blob.bucket.name, test_blob.name)
     out, _ = capsys.readouterr()
diff --git a/samples/snippets/storage_download_to_stream.py b/samples/snippets/storage_download_to_stream.py
@@ -0,0 +1,50 @@
+#!/usr/bin/env python
+
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# [START storage_stream_file_download]
+from google.cloud import storage
+
+
+def download_blob_to_stream(bucket_name, source_blob_name, file_obj):
+    """Downloads a blob to a stream or other file-like object."""
+
+    # The ID of your GCS bucket
+    # bucket_name = "your-bucket-name"
+
+    # The ID of your GCS object (blob)
+    # source_blob_name = "storage-object-name"
+
+    # The stream or file (file-like object) to which the blob will be written
+    # import io
+    # file_obj = io.BytesIO()
+
+    storage_client = storage.Client()
+
+    bucket = storage_client.bucket(bucket_name)
+
+    # Construct a client-side representation of a blob.
+    # Note `Bucket.blob` differs from `Bucket.get_blob` in that it doesn't
+    # retrieve metadata from Google Cloud Storage. As we don't use metadata in
+    # this example, using `Bucket.blob` is preferred here.
+    blob = bucket.blob(source_blob_name)
+    blob.download_to_file(file_obj)
+
+    print("Downloaded blob {} to file-like object.".format(source_blob_name))
+
+    return file_obj
+    # Before reading from file_obj, remember to rewind with file_obj.seek(0).
+
+# [END storage_stream_file_download]
diff --git a/samples/snippets/storage_upload_from_memory.py b/samples/snippets/storage_upload_from_memory.py
@@ -22,10 +22,13 @@
 
 def upload_blob_from_memory(bucket_name, contents, destination_blob_name):
     """Uploads a file to the bucket."""
+
     # The ID of your GCS bucket
     # bucket_name = "your-bucket-name"
+
     # The contents to upload to the file
     # contents = "these are my contents"
+
     # The ID of your GCS object
     # destination_blob_name = "storage-object-name"
 
@@ -37,13 +40,13 @@ def upload_blob_from_memory(bucket_name, contents, destination_blob_name):
 
     print(
         "{} with contents {} uploaded to {}.".format(
-            destination_blob_name, contents, destination_blob_name
+            destination_blob_name, contents, bucket_name
         )
     )
 
-
 # [END storage_file_upload_from_memory]
 
+
 if __name__ == "__main__":
     upload_blob_from_memory(
         bucket_name=sys.argv[1],
diff --git a/samples/snippets/storage_upload_from_stream.py b/samples/snippets/storage_upload_from_stream.py
@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# [START storage_stream_file_upload]
+from google.cloud import storage
+
+
+def upload_blob_from_stream(bucket_name, file_obj, destination_blob_name):
+    """Uploads bytes from a stream or other file-like object to a blob."""
+    # The ID of your GCS bucket
+    # bucket_name = "your-bucket-name"
+
+    # The stream or file (file-like object) from which to read
+    # import io
+    # file_obj = io.StringIO()
+    # file_obj.write("This is test data.")
+
+    # The desired name of the uploaded GCS object (blob)
+    # destination_blob_name = "storage-object-name"
+
+    # Construct a client-side representation of the blob.
+    storage_client = storage.Client()
+    bucket = storage_client.bucket(bucket_name)
+    blob = bucket.blob(destination_blob_name)
+
+    # Rewind the stream to the beginning. This step can be omitted if the input
+    # stream will always be at a correct position.
+    file_obj.seek(0)
+
+    # Upload data from the stream to your bucket.
+    blob.upload_from_file(file_obj)
+
+    print(
+        "Stream data uploaded to {} in bucket {}.".format(
+            destination_blob_name, bucket_name
+        )
+    )
+
+# [END storage_stream_file_upload]
