From 88c802c9d3e9ade05468e0e5a574cbbe89bae9c4 Mon Sep 17 00:00:00 2001
From: Charles Chen <ccy@google.com>
Date: Mon, 22 Aug 2016 16:16:51 -0700
Subject: [PATCH] Support mode attribute on GCS files

---
 sdks/python/apache_beam/io/gcsio.py      | 13 +++++++++----
 sdks/python/apache_beam/io/gcsio_test.py | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/sdks/python/apache_beam/io/gcsio.py b/sdks/python/apache_beam/io/gcsio.py
index 4c733d9b178e..6d02a251b3e1 100644
--- a/sdks/python/apache_beam/io/gcsio.py
+++ b/sdks/python/apache_beam/io/gcsio.py
@@ -111,10 +111,11 @@ def open(self, filename, mode='r',
       ValueError: Invalid open file mode.
     """
     if mode == 'r' or mode == 'rb':
-      return GcsBufferedReader(self.client, filename,
+      return GcsBufferedReader(self.client, filename, mode=mode,
                                buffer_size=read_buffer_size)
     elif mode == 'w' or mode == 'wb':
-      return GcsBufferedWriter(self.client, filename, mime_type=mime_type)
+      return GcsBufferedWriter(self.client, filename, mode=mode,
+                               mime_type=mime_type)
     else:
       raise ValueError('Invalid file open mode: %s.' % mode)
 
@@ -262,11 +263,13 @@ def size(self, path):
 class GcsBufferedReader(object):
   """A class for reading Google Cloud Storage files."""
 
-  def __init__(self, client, path, buffer_size=DEFAULT_READ_BUFFER_SIZE):
+  def __init__(self, client, path, mode='r',
+               buffer_size=DEFAULT_READ_BUFFER_SIZE):
     self.client = client
     self.path = path
     self.bucket, self.name = parse_gcs_path(path)
     self.buffer_size = buffer_size
+    self.mode = mode
 
     # Get object state.
     get_request = (
@@ -545,10 +548,12 @@ def _check_open(self):
       if self.closed:
         raise IOError('Stream is closed.')
 
-  def __init__(self, client, path, mime_type='application/octet-stream'):
+  def __init__(self, client, path, mode='w',
+               mime_type='application/octet-stream'):
     self.client = client
     self.path = path
     self.bucket, self.name = parse_gcs_path(path)
+    self.mode = mode
 
     self.closed = False
     self.position = 0
diff --git a/sdks/python/apache_beam/io/gcsio_test.py b/sdks/python/apache_beam/io/gcsio_test.py
index 99c99b38201c..1e2c50e791d6 100644
--- a/sdks/python/apache_beam/io/gcsio_test.py
+++ b/sdks/python/apache_beam/io/gcsio_test.py
@@ -218,6 +218,20 @@ def test_size(self):
     self.assertTrue(self.gcs.exists(file_name))
     self.assertEqual(1234, self.gcs.size(file_name))
 
+  def test_file_mode(self):
+    file_name = 'gs://gcsio-test/dummy_mode_file'
+    with self.gcs.open(file_name, 'wb') as f:
+      assert f.mode == 'wb'
+    with self.gcs.open(file_name, 'rb') as f:
+      assert f.mode == 'rb'
+
+  def test_bad_file_modes(self):
+    file_name = 'gs://gcsio-test/dummy_mode_file'
+    with self.assertRaises(ValueError):
+      self.gcs.open(file_name, 'w+')
+    with self.assertRaises(ValueError):
+      self.gcs.open(file_name, 'r+b')
+
   def test_delete(self):
     file_name = 'gs://gcsio-test/delete_me'
     file_size = 1024