Propagate the tensorflow::Status as a Python exception for various file open

operations.

Note that this changes the exception type of file not found errors from IOError to errors.NotFoundError.
Change: 133720319

11 files changed, 57 insertions, 31 deletions

diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index e344f6be70..01996d67ae 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -140,6 +140,8 @@ cc_library(
         "lib/io/py_record_reader.h",
     ],
     deps = [
+        "//tensorflow/c:c_api",
+        "//tensorflow/c:tf_status_helper",
         "//tensorflow/core:lib",
     ],
 )
@@ -153,6 +155,8 @@ cc_library(
         "lib/io/py_record_writer.h",
     ],
     deps = [
+        "//tensorflow/c:c_api",
+        "//tensorflow/c:tf_status_helper",
         "//tensorflow/core:lib",
     ],
 )
diff --git a/tensorflow/python/client/events_writer_test.py b/tensorflow/python/client/events_writer_test.py
index c0be53899e..f02a83cbde 100644
--- a/tensorflow/python/client/events_writer_test.py
+++ b/tensorflow/python/client/events_writer_test.py
@@ -24,6 +24,7 @@ from tensorflow.core.framework import summary_pb2
 from tensorflow.core.util import event_pb2
 from tensorflow.python import pywrap_tensorflow
 from tensorflow.python.lib.io import tf_record
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.platform import googletest
 from tensorflow.python.util import compat
@@ -43,7 +44,7 @@ class PywrapeventsWriterTest(test_util.TensorFlowTestCase):
     writer.Flush()
     writer.Close()
 
-    with self.assertRaises(IOError):
+    with self.assertRaises(errors.NotFoundError):
       for r in tf_record.tf_record_iterator(filename + "DOES_NOT_EXIST"):
         self.assertTrue(False)
 
diff --git a/tensorflow/python/lib/io/file_io.i b/tensorflow/python/lib/io/file_io.i
index 8bfba4b90b..e4a9aa94c5 100644
--- a/tensorflow/python/lib/io/file_io.i
+++ b/tensorflow/python/lib/io/file_io.i
@@ -169,9 +169,12 @@ void Stat(const string& filename, FileStatistics* stats,
 }
 
 tensorflow::io::BufferedInputStream* CreateBufferedInputStream(
-    const string& filename, size_t buffer_size) {
+    const string& filename, size_t buffer_size, TF_Status* out_status) {
   std::unique_ptr<tensorflow::RandomAccessFile> file;
-  if (!tensorflow::Env::Default()->NewRandomAccessFile(filename, &file).ok()) {
+  tensorflow::Status status =
+      tensorflow::Env::Default()->NewRandomAccessFile(filename, &file);
+  if (!status.ok()) {
+    Set_TF_Status_from_Status(out_status, status);
     return nullptr;
   }
   std::unique_ptr<tensorflow::io::RandomAccessInputStream> input_stream(
@@ -182,9 +185,13 @@ tensorflow::io::BufferedInputStream* CreateBufferedInputStream(
   return buffered_input_stream.release();
 }
 
-tensorflow::WritableFile* CreateWritableFile(const string& filename) {
+tensorflow::WritableFile* CreateWritableFile(
+    const string& filename, TF_Status* out_status) {
   std::unique_ptr<tensorflow::WritableFile> file;
-  if (!tensorflow::Env::Default()->NewWritableFile(filename, &file).ok()) {
+  tensorflow::Status status =
+      tensorflow::Env::Default()->NewWritableFile(filename, &file);
+  if (!status.ok()) {
+    Set_TF_Status_from_Status(out_status, status);
     return nullptr;
   }
   return file.release();
@@ -230,8 +237,9 @@ bool IsDirectory(const string& dirname, TF_Status* out_status);
 void Stat(const string& filename, tensorflow::FileStatistics* stats,
           TF_Status* out_status);
 tensorflow::io::BufferedInputStream* CreateBufferedInputStream(
-    const string& filename, size_t buffer_size);
-tensorflow::WritableFile* CreateWritableFile(const string& filename);
+    const string& filename, size_t buffer_size, TF_Status* out_status);
+tensorflow::WritableFile* CreateWritableFile(const string& filename,
+                                             TF_Status* out_status);
 void AppendToFile(const string& file_content, tensorflow::WritableFile* file,
                   TF_Status* out_status);
 void FlushWritableFile(tensorflow::WritableFile* file, TF_Status* out_status);
diff --git a/tensorflow/python/lib/io/file_io.py b/tensorflow/python/lib/io/file_io.py
index ba1e18bae6..f2fede278b 100644
--- a/tensorflow/python/lib/io/file_io.py
+++ b/tensorflow/python/lib/io/file_io.py
@@ -67,22 +67,18 @@ class FileIO(object):
       if not self._read_check_passed:
         raise errors.PermissionDeniedError(None, None,
                                            "File isn't open for reading")
-      self._read_buf = pywrap_tensorflow.CreateBufferedInputStream(
-          compat.as_bytes(self.__name), 1024 * 512)
-      if not self._read_buf:
-        raise errors.InternalError(None, None,
-                                   "Could not open file for streaming")
+      with errors.raise_exception_on_not_ok_status() as status:
+        self._read_buf = pywrap_tensorflow.CreateBufferedInputStream(
+            compat.as_bytes(self.__name), 1024 * 512, status)
 
   def _prewrite_check(self):
     if not self._writable_file:
       if not self._write_check_passed:
         raise errors.PermissionDeniedError(None, None,
                                            "File isn't open for writing")
-      self._writable_file = pywrap_tensorflow.CreateWritableFile(
-          compat.as_bytes(self.__name))
-      if not self._writable_file:
-        raise errors.InternalError(None, None,
-                                   "Could not open file for writing")
+      with errors.raise_exception_on_not_ok_status() as status:
+        self._writable_file = pywrap_tensorflow.CreateWritableFile(
+            compat.as_bytes(self.__name), status)
 
   def size(self):
     """Returns the size of the file."""
diff --git a/tensorflow/python/lib/io/py_record_reader.cc b/tensorflow/python/lib/io/py_record_reader.cc
index 552cca6a4b..47c0878932 100644
--- a/tensorflow/python/lib/io/py_record_reader.cc
+++ b/tensorflow/python/lib/io/py_record_reader.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/python/lib/io/py_record_reader.h"
 
+#include "tensorflow/c/tf_status_helper.h"
 #include "tensorflow/core/lib/core/stringpiece.h"
 #include "tensorflow/core/lib/io/record_reader.h"
 #include "tensorflow/core/lib/io/zlib_compression_options.h"
@@ -30,10 +31,12 @@ namespace io {
 PyRecordReader::PyRecordReader() {}
 
 PyRecordReader* PyRecordReader::New(const string& filename, uint64 start_offset,
-                                    const string& compression_type_string) {
+                                    const string& compression_type_string,
+                                    TF_Status* out_status) {
   std::unique_ptr<RandomAccessFile> file;
   Status s = Env::Default()->NewRandomAccessFile(filename, &file);
   if (!s.ok()) {
+    Set_TF_Status_from_Status(out_status, s);
     return nullptr;
   }
   PyRecordReader* reader = new PyRecordReader;
diff --git a/tensorflow/python/lib/io/py_record_reader.h b/tensorflow/python/lib/io/py_record_reader.h
index a72cf04c3c..0da74ee948 100644
--- a/tensorflow/python/lib/io/py_record_reader.h
+++ b/tensorflow/python/lib/io/py_record_reader.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_PYTHON_LIB_IO_PY_RECORD_READER_H_
 #define TENSORFLOW_PYTHON_LIB_IO_PY_RECORD_READER_H_
 
+#include "tensorflow/c/c_api.h"
 #include "tensorflow/core/lib/core/stringpiece.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/types.h"
@@ -36,7 +37,8 @@ class PyRecordReader {
   // TODO(vrv): make this take a shared proto to configure
   // the compression options.
   static PyRecordReader* New(const string& filename, uint64 start_offset,
-                             const string& compression_type_string);
+                             const string& compression_type_string,
+                             TF_Status* out_status);
 
   ~PyRecordReader();
 
diff --git a/tensorflow/python/lib/io/py_record_writer.cc b/tensorflow/python/lib/io/py_record_writer.cc
index 019f5e2fac..d9fdda7ebf 100644
--- a/tensorflow/python/lib/io/py_record_writer.cc
+++ b/tensorflow/python/lib/io/py_record_writer.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/python/lib/io/py_record_writer.h"
 
+#include "tensorflow/c/tf_status_helper.h"
 #include "tensorflow/core/lib/core/stringpiece.h"
 #include "tensorflow/core/lib/io/record_writer.h"
 #include "tensorflow/core/lib/io/zlib_compression_options.h"
@@ -27,10 +28,12 @@ namespace io {
 PyRecordWriter::PyRecordWriter() {}
 
 PyRecordWriter* PyRecordWriter::New(const string& filename,
-                                    const string& compression_type_string) {
+                                    const string& compression_type_string,
+                                    TF_Status* out_status) {
   std::unique_ptr<WritableFile> file;
   Status s = Env::Default()->NewWritableFile(filename, &file);
   if (!s.ok()) {
+    Set_TF_Status_from_Status(out_status, s);
     return nullptr;
   }
   PyRecordWriter* writer = new PyRecordWriter;
diff --git a/tensorflow/python/lib/io/py_record_writer.h b/tensorflow/python/lib/io/py_record_writer.h
index 86e2b9e56f..483b7b9df0 100644
--- a/tensorflow/python/lib/io/py_record_writer.h
+++ b/tensorflow/python/lib/io/py_record_writer.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_PYTHON_LIB_IO_PY_RECORD_WRITER_H_
 #define TENSORFLOW_PYTHON_LIB_IO_PY_RECORD_WRITER_H_
 
+#include "tensorflow/c/c_api.h"
 #include "tensorflow/core/lib/core/stringpiece.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/types.h"
@@ -36,7 +37,8 @@ class PyRecordWriter {
   // TODO(vrv): make this take a shared proto to configure
   // the compression options.
   static PyRecordWriter* New(const string& filename,
-                             const string& compression_type_string);
+                             const string& compression_type_string,
+                             TF_Status* out_status);
   ~PyRecordWriter();
 
   bool WriteRecord(tensorflow::StringPiece record);
diff --git a/tensorflow/python/lib/io/tf_record.py b/tensorflow/python/lib/io/tf_record.py
index 2419a1ef1a..96b212c8ad 100644
--- a/tensorflow/python/lib/io/tf_record.py
+++ b/tensorflow/python/lib/io/tf_record.py
@@ -20,6 +20,7 @@ from __future__ import division
 from __future__ import print_function
 
 from tensorflow.python import pywrap_tensorflow
+from tensorflow.python.framework import errors
 from tensorflow.python.util import compat
 
 
@@ -59,8 +60,10 @@ def tf_record_iterator(path, options=None):
     IOError: If `path` cannot be opened for reading.
   """
   compression_type_string = options.get_type_as_string() if options else ""
-  reader = pywrap_tensorflow.PyRecordReader_New(
-      compat.as_bytes(path), 0, compat.as_bytes(compression_type_string))
+  with errors.raise_exception_on_not_ok_status() as status:
+    reader = pywrap_tensorflow.PyRecordReader_New(
+        compat.as_bytes(path), 0, compat.as_bytes(compression_type_string),
+        status)
 
   if reader is None:
     raise IOError("Could not open %s." % path)
@@ -93,10 +96,10 @@ class TFRecordWriter(object):
     """
     compression_type_string = options.get_type_as_string() if options else ""
 
-    self._writer = pywrap_tensorflow.PyRecordWriter_New(
-        compat.as_bytes(path), compat.as_bytes(compression_type_string))
-    if self._writer is None:
-      raise IOError("Could not write to %s." % path)
+    with errors.raise_exception_on_not_ok_status() as status:
+      self._writer = pywrap_tensorflow.PyRecordWriter_New(
+          compat.as_bytes(path), compat.as_bytes(compression_type_string),
+          status)
 
   def __enter__(self):
     """Enter a `with` block."""
diff --git a/tensorflow/python/summary/impl/event_file_loader.py b/tensorflow/python/summary/impl/event_file_loader.py
index 509a371412..dedebe5484 100644
--- a/tensorflow/python/summary/impl/event_file_loader.py
+++ b/tensorflow/python/summary/impl/event_file_loader.py
@@ -20,6 +20,7 @@ from __future__ import print_function
 
 from tensorflow.core.util import event_pb2
 from tensorflow.python import pywrap_tensorflow
+from tensorflow.python.framework import errors
 from tensorflow.python.platform import app
 from tensorflow.python.platform import resource_loader
 from tensorflow.python.platform import tf_logging as logging
@@ -34,8 +35,9 @@ class EventFileLoader(object):
       raise ValueError('A file path is required')
     file_path = resource_loader.readahead_file_path(file_path)
     logging.debug('Opening a record reader pointing at %s', file_path)
-    self._reader = pywrap_tensorflow.PyRecordReader_New(
-        compat.as_bytes(file_path), 0, compat.as_bytes(''))
+    with errors.raise_exception_on_not_ok_status() as status:
+      self._reader = pywrap_tensorflow.PyRecordReader_New(
+          compat.as_bytes(file_path), 0, compat.as_bytes(''), status)
     # Store it for logging purposes.
     self._file_path = file_path
     if not self._reader:
diff --git a/tensorflow/python/summary/impl/gcs_file_loader.py b/tensorflow/python/summary/impl/gcs_file_loader.py
index b82bb85fc2..c46534dbb5 100644
--- a/tensorflow/python/summary/impl/gcs_file_loader.py
+++ b/tensorflow/python/summary/impl/gcs_file_loader.py
@@ -21,6 +21,7 @@ import tempfile
 
 from tensorflow.core.util import event_pb2
 from tensorflow.python import pywrap_tensorflow
+from tensorflow.python.framework import errors
 from tensorflow.python.platform import app
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.summary.impl import gcs
@@ -46,8 +47,9 @@ class GCSFileLoader(object):
       name = temp_file.name
       logging.debug('Temp file created at %s', name)
       gcs.CopyContents(self._gcs_path, self._gcs_offset, temp_file)
-      reader = pywrap_tensorflow.PyRecordReader_New(
-          compat.as_bytes(name), 0, compat.as_bytes(''))
+      with errors.raise_exception_on_not_ok_status() as status:
+        reader = pywrap_tensorflow.PyRecordReader_New(
+            compat.as_bytes(name), 0, compat.as_bytes(''), status)
       while reader.GetNext():
         event = event_pb2.Event()
         event.ParseFromString(reader.record())