aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/python/lib
diff options
context:
space:
mode:
Diffstat (limited to 'tensorflow/python/lib')
-rwxr-xr-xtensorflow/python/lib/__init__.py0
-rwxr-xr-xtensorflow/python/lib/core/__init__.py0
-rw-r--r--tensorflow/python/lib/core/pywrap_status_test.py35
-rw-r--r--tensorflow/python/lib/core/status.i116
-rw-r--r--tensorflow/python/lib/core/status_helper.i16
-rw-r--r--tensorflow/python/lib/core/strings.i94
-rwxr-xr-xtensorflow/python/lib/io/__init__.py0
-rw-r--r--tensorflow/python/lib/io/py_record_reader.cc49
-rw-r--r--tensorflow/python/lib/io/py_record_reader.h50
-rw-r--r--tensorflow/python/lib/io/py_record_reader.i39
-rw-r--r--tensorflow/python/lib/io/py_record_writer.cc44
-rw-r--r--tensorflow/python/lib/io/py_record_writer.h38
-rw-r--r--tensorflow/python/lib/io/py_record_writer.i38
-rw-r--r--tensorflow/python/lib/io/python_io.py29
-rw-r--r--tensorflow/python/lib/io/tf_record.py68
15 files changed, 616 insertions, 0 deletions
diff --git a/tensorflow/python/lib/__init__.py b/tensorflow/python/lib/__init__.py
new file mode 100755
index 0000000000..e69de29bb2
--- /dev/null
+++ b/tensorflow/python/lib/__init__.py
diff --git a/tensorflow/python/lib/core/__init__.py b/tensorflow/python/lib/core/__init__.py
new file mode 100755
index 0000000000..e69de29bb2
--- /dev/null
+++ b/tensorflow/python/lib/core/__init__.py
diff --git a/tensorflow/python/lib/core/pywrap_status_test.py b/tensorflow/python/lib/core/pywrap_status_test.py
new file mode 100644
index 0000000000..000a784b6c
--- /dev/null
+++ b/tensorflow/python/lib/core/pywrap_status_test.py
@@ -0,0 +1,35 @@
+"""Tests for SWIG wrapped brain::Status."""
+
+from tensorflow.core.lib.core import error_codes_pb2
+from tensorflow.python import pywrap_tensorflow
+from tensorflow.python.platform import googletest
+
+
+class StatusTest(googletest.TestCase):
+
+ def testDefaultOk(self):
+ status = pywrap_tensorflow.Status()
+ self.assertTrue(status.ok())
+
+ def testCodeAndMessage(self):
+ status = pywrap_tensorflow.Status(error_codes_pb2.INVALID_ARGUMENT, 'foo')
+ self.assertEqual(error_codes_pb2.INVALID_ARGUMENT, status.code())
+ self.assertEqual('foo', status.error_message())
+
+ def testToString(self):
+ status = pywrap_tensorflow.Status()
+ # .ToString was remapped in the .swig file, hence will not work
+ # self.assertIn('OK', status.ToString())
+ self.assertIn('OK', str(status))
+
+ def testException(self):
+ with self.assertRaises(pywrap_tensorflow.StatusNotOK) as context:
+ pywrap_tensorflow.NotOkay()
+ self.assertEqual(context.exception.code, error_codes_pb2.INVALID_ARGUMENT)
+ self.assertEqual(context.exception.error_message, 'Testing 1 2 3')
+ self.assertEqual(None, pywrap_tensorflow.Okay(),
+ 'Status wrapper should not return anything upon OK.')
+
+
+if __name__ == '__main__':
+ googletest.main()
diff --git a/tensorflow/python/lib/core/status.i b/tensorflow/python/lib/core/status.i
new file mode 100644
index 0000000000..fddbc31e24
--- /dev/null
+++ b/tensorflow/python/lib/core/status.i
@@ -0,0 +1,116 @@
+// SWIG wrapper for lib::tensorflow::Status
+
+%include "tensorflow/python/platform/base.i"
+%include "tensorflow/python/lib/core/strings.i"
+
+%apply int { tensorflow::error::Code }; // Treat the enum as an integer.
+
+%{
+#include "tensorflow/core/public/status.h"
+%}
+
+%typemap(out, fragment="StatusNotOK") tensorflow::Status {
+ if ($1.ok()) {
+ $result = SWIG_Py_Void();
+ } else {
+ RaiseStatusNotOK($1, $descriptor(tensorflow::Status*));
+ SWIG_fail;
+ }
+}
+
+%init %{
+// Setup the StatusNotOK exception class.
+PyObject *pywrap_status = PyImport_ImportModuleNoBlock(
+ "tensorflow.python.pywrap_tensorflow");
+if (pywrap_status) {
+ PyObject *exception = PyErr_NewException(
+ "tensorflow.python.pywrap_tensorflow.StatusNotOK",
+ NULL, NULL);
+ if (exception) {
+ PyModule_AddObject(pywrap_status, "StatusNotOK", exception); // Steals ref.
+ }
+ Py_DECREF(pywrap_status);
+}
+%}
+
+%fragment("StatusNotOK", "header") %{
+#include "tensorflow/core/public/status.h"
+
+namespace {
+// Initialized on the first call to RaiseStatusNotOK().
+static PyObject *StatusNotOKError = nullptr;
+
+inline void Py_DECREF_wrapper(PyObject *o) { Py_DECREF(o); }
+typedef std::unique_ptr<PyObject, decltype(&Py_DECREF_wrapper)> SafePyObjectPtr;
+SafePyObjectPtr make_safe(PyObject* o) {
+ return SafePyObjectPtr(o, Py_DECREF_wrapper);
+}
+
+void RaiseStatusNotOK(const tensorflow::Status& status, swig_type_info *type) {
+ const int code = status.code();
+ string fullmsg = status.ToString();
+
+ PyObject *exception = nullptr;
+
+ // We're holding the Python GIL, so we don't need to synchronize
+ // access to StatusNotOKError with a Mutex of our own.
+ if (!StatusNotOKError) {
+ PyObject *cls = nullptr;
+ auto pywrap = make_safe(PyImport_ImportModule(
+ "tensorflow.python.pywrap_tensorflow"));
+ if (pywrap) {
+ cls = PyObject_GetAttrString(pywrap.get(), "StatusNotOK");
+ }
+ if (!cls) {
+ cls = Py_None;
+ Py_INCREF(cls);
+ }
+ StatusNotOKError = cls;
+ }
+
+ if (StatusNotOKError != Py_None) {
+ auto fullmsg_ptr = make_safe(_SwigString_FromString(fullmsg));
+ auto exception_ptr = make_safe(PyObject_CallFunctionObjArgs(
+ StatusNotOKError, fullmsg_ptr.get(), NULL));
+ exception = exception_ptr.get();
+ if (exception) {
+ auto pycode = make_safe(PyInt_FromLong(static_cast<long>(code)));
+ auto pymsg = make_safe(_SwigString_FromString(status.error_message()));
+ auto pystatus = make_safe(SWIG_NewPointerObj(
+ SWIG_as_voidptr(new tensorflow::Status(status)), type, SWIG_POINTER_OWN));
+ PyObject_SetAttrString(exception, "code", pycode.get());
+ PyObject_SetAttrString(exception, "error_message", pymsg.get());
+ PyErr_SetObject(StatusNotOKError, exception);
+ }
+ }
+ if (!exception) {
+ fullmsg =
+ ("could not construct StatusNotOK (original error "
+ " was: " +
+ fullmsg + ")");
+ PyErr_SetString(PyExc_SystemError, fullmsg.c_str());
+ }
+}
+
+} // namespace
+%}
+
+%ignoreall
+
+%unignore tensorflow;
+%unignore tensorflow::lib;
+%unignore tensorflow::Status;
+%unignore tensorflow::Status::Status;
+%unignore tensorflow::Status::Status(tensorflow::error::Code, StringPiece);
+%unignore tensorflow::Status::~Status;
+%unignore tensorflow::Status::code;
+%unignore tensorflow::Status::ok;
+%unignore tensorflow::Status::error_message;
+%unignore tensorflow::Status::ToString;
+%ignore tensorflow::Status::operator=;
+
+%rename(__str__) tensorflow::Status::ToString;
+
+%include "tensorflow/core/public/status.h"
+
+%unignoreall
diff --git a/tensorflow/python/lib/core/status_helper.i b/tensorflow/python/lib/core/status_helper.i
new file mode 100644
index 0000000000..2e01e79ebd
--- /dev/null
+++ b/tensorflow/python/lib/core/status_helper.i
@@ -0,0 +1,16 @@
+// SWIG test helper for lib::tensorflow::Status
+
+%include "tensorflow/python/platform/base.i"
+%import(module="tensorflow.python.pywrap_tensorflow") "tensorflow/python/lib/core/status.i"
+
+%inline %{
+#include "tensorflow/core/public/status.h"
+
+tensorflow::Status NotOkay() {
+ return tensorflow::Status(tensorflow::error::INVALID_ARGUMENT, "Testing 1 2 3");
+}
+
+tensorflow::Status Okay() {
+ return tensorflow::Status();
+}
+%}
diff --git a/tensorflow/python/lib/core/strings.i b/tensorflow/python/lib/core/strings.i
new file mode 100644
index 0000000000..c88e426a54
--- /dev/null
+++ b/tensorflow/python/lib/core/strings.i
@@ -0,0 +1,94 @@
+// Wrapper functions to provide a scripting-language-friendly interface
+// to our string libraries.
+//
+// NOTE: as of 2005-01-13, this SWIG file is not used to generate a pywrap
+// library for manipulation of various string-related types or access
+// to the special string functions (Python has plenty). This SWIG file
+// should be %import'd so that other SWIG wrappers have proper access
+// to the types in //strings (such as the StringPiece object). We may
+// generate a pywrap at some point in the future.
+//
+// NOTE: (Dan Ardelean) as of 2005-11-15 added typemaps to convert Java String
+// arguments to C++ StringPiece& objects. This is required because a
+// StringPiece class does not make sense - the code SWIG generates for a
+// StringPiece class is useless, because it releases the buffer set in
+// StringPiece after creating the object. C++ StringPiece objects rely on
+// the buffer holding the data being allocated externally.
+
+// NOTE: for now, we'll just start with what is needed, and add stuff
+// as it comes up.
+
+%{
+#include "tensorflow/core/lib/core/stringpiece.h"
+%}
+
+%typemap(typecheck) tensorflow::StringPiece = char *;
+%typemap(typecheck) const tensorflow::StringPiece & = char *;
+
+// "tensorflow::StringPiece" arguments can be provided by a simple Python 'str' string
+// or a 'unicode' object. If 'unicode', it's translated using the default
+// encoding, i.e., sys.getdefaultencoding(). If passed None, a tensorflow::StringPiece
+// of zero length with a NULL pointer is provided.
+%typemap(in) tensorflow::StringPiece {
+ if ($input != Py_None) {
+ char * buf;
+ Py_ssize_t len;
+%#if PY_VERSION_HEX >= 0x03030000
+ /* Do unicode handling as PyBytes_AsStringAndSize doesn't in Python 3. */
+ if (PyUnicode_Check($input)) {
+ buf = PyUnicode_AsUTF8AndSize($input, &len);
+ if (buf == NULL)
+ SWIG_fail;
+ } else {
+%#elif PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION < 3
+%# error "Unsupported Python 3.x C API version (3.3 or later required)."
+%#endif
+ if (PyBytes_AsStringAndSize($input, &buf, &len) == -1) {
+ // Python has raised an error (likely TypeError or UnicodeEncodeError).
+ SWIG_fail;
+ }
+%#if PY_VERSION_HEX >= 0x03030000
+ }
+%#endif
+ $1.set(buf, len);
+ }
+}
+
+// "const tensorflow::StringPiece&" arguments can be provided the same as
+// "tensorflow::StringPiece", whose typemap is defined above.
+%typemap(in) const tensorflow::StringPiece & (tensorflow::StringPiece temp) {
+ if ($input != Py_None) {
+ char * buf;
+ Py_ssize_t len;
+%#if PY_VERSION_HEX >= 0x03030000
+ /* Do unicode handling as PyBytes_AsStringAndSize doesn't in Python 3. */
+ if (PyUnicode_Check($input)) {
+ buf = PyUnicode_AsUTF8AndSize($input, &len);
+ if (buf == NULL)
+ SWIG_fail;
+ } else {
+%#elif PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION < 3
+%# error "Unsupported Python 3.x C API version (3.3 or later required)."
+%#endif
+ if (PyBytes_AsStringAndSize($input, &buf, &len) == -1) {
+ // Python has raised an error (likely TypeError or UnicodeEncodeError).
+ SWIG_fail;
+ }
+%#if PY_VERSION_HEX >= 0x03030000
+ }
+%#endif
+ temp.set(buf, len);
+ }
+ $1 = &temp;
+}
+
+// C++ functions returning tensorflow::StringPiece will simply return bytes in Python,
+// or None if the StringPiece contained a NULL pointer.
+%typemap(out) tensorflow::StringPiece {
+ if ($1.data()) {
+ $result = PyString_FromStringAndSize($1.data(), $1.size());
+ } else {
+ Py_INCREF(Py_None);
+ $result = Py_None;
+ }
+}
diff --git a/tensorflow/python/lib/io/__init__.py b/tensorflow/python/lib/io/__init__.py
new file mode 100755
index 0000000000..e69de29bb2
--- /dev/null
+++ b/tensorflow/python/lib/io/__init__.py
diff --git a/tensorflow/python/lib/io/py_record_reader.cc b/tensorflow/python/lib/io/py_record_reader.cc
new file mode 100644
index 0000000000..5cc5229a8b
--- /dev/null
+++ b/tensorflow/python/lib/io/py_record_reader.cc
@@ -0,0 +1,49 @@
+#include "tensorflow/python/lib/io/py_record_reader.h"
+
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/io/record_reader.h"
+#include "tensorflow/core/platform/port.h"
+#include "tensorflow/core/public/env.h"
+
+namespace tensorflow {
+
+class RandomAccessFile;
+
+namespace io {
+
+PyRecordReader::PyRecordReader() {}
+
+PyRecordReader* PyRecordReader::New(const string& filename,
+ uint64 start_offset) {
+ RandomAccessFile* file;
+ Status s = Env::Default()->NewRandomAccessFile(filename, &file);
+ if (!s.ok()) {
+ return nullptr;
+ }
+ PyRecordReader* reader = new PyRecordReader;
+ reader->offset_ = start_offset;
+ reader->file_ = file;
+ reader->reader_ = new RecordReader(reader->file_);
+ return reader;
+}
+
+PyRecordReader::~PyRecordReader() {
+ delete reader_;
+ delete file_;
+}
+
+bool PyRecordReader::GetNext() {
+ if (reader_ == nullptr) return false;
+ Status s = reader_->ReadRecord(&offset_, &record_);
+ return s.ok();
+}
+
+void PyRecordReader::Close() {
+ delete reader_;
+ delete file_;
+ file_ = nullptr;
+ reader_ = nullptr;
+}
+
+} // namespace io
+} // namespace tensorflow
diff --git a/tensorflow/python/lib/io/py_record_reader.h b/tensorflow/python/lib/io/py_record_reader.h
new file mode 100644
index 0000000000..5a775761df
--- /dev/null
+++ b/tensorflow/python/lib/io/py_record_reader.h
@@ -0,0 +1,50 @@
+#ifndef TENSORFLOW_PYTHON_LIB_IO_PY_RECORD_READER_H_
+#define TENSORFLOW_PYTHON_LIB_IO_PY_RECORD_READER_H_
+
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/platform/port.h"
+#include "tensorflow/core/public/status.h"
+
+namespace tensorflow {
+
+class RandomAccessFile;
+
+namespace io {
+
+class RecordReader;
+
+// A wrapper around io::RecordReader that is more easily SWIG wrapped for
+// Python. An instance of this class is not safe for concurrent access
+// by multiple threads.
+class PyRecordReader {
+ public:
+ static PyRecordReader* New(const string& filename, uint64 start_offset);
+ ~PyRecordReader();
+
+ // Attempt to get the next record at "current_offset()". If
+ // successful, returns true, and the record contents can be retrieve
+ // with "this->record()". Otherwise, returns false.
+ bool GetNext();
+ // Return the current record contents. Only valid after the preceding call
+ // to GetNext() returned true
+ string record() const { return record_; }
+ // Return the current offset in the file.
+ uint64 offset() const { return offset_; }
+
+ // Close the underlying file and release its resources.
+ void Close();
+
+ private:
+ PyRecordReader();
+
+ uint64 offset_;
+ RandomAccessFile* file_; // Owned
+ io::RecordReader* reader_; // Owned
+ string record_;
+ TF_DISALLOW_COPY_AND_ASSIGN(PyRecordReader);
+};
+
+} // namespace io
+} // namespace tensorflow
+
+#endif // TENSORFLOW_PYTHON_LIB_IO_PY_RECORD_READER_H_
diff --git a/tensorflow/python/lib/io/py_record_reader.i b/tensorflow/python/lib/io/py_record_reader.i
new file mode 100644
index 0000000000..19f911bd52
--- /dev/null
+++ b/tensorflow/python/lib/io/py_record_reader.i
@@ -0,0 +1,39 @@
+%nothread tensorflow::io::PyRecordReader::GetNext;
+
+%include "tensorflow/python/platform/base.i"
+
+%feature("except") tensorflow::io::PyRecordReader::New {
+ // Let other threads run while we read
+ Py_BEGIN_ALLOW_THREADS
+ $action
+ Py_END_ALLOW_THREADS
+}
+
+%newobject tensorflow::io::PyRecordReader::New;
+
+%feature("except") tensorflow::io::PyRecordReader::GetNext {
+ // Let other threads run while we read
+ Py_BEGIN_ALLOW_THREADS
+ $action
+ Py_END_ALLOW_THREADS
+}
+
+%{
+#include "tensorflow/python/lib/io/py_record_reader.h"
+%}
+
+%ignoreall
+
+%unignore tensorflow;
+%unignore tensorflow::io;
+%unignore tensorflow::io::PyRecordReader;
+%unignore tensorflow::io::PyRecordReader::~PyRecordReader;
+%unignore tensorflow::io::PyRecordReader::GetNext;
+%unignore tensorflow::io::PyRecordReader::offset;
+%unignore tensorflow::io::PyRecordReader::record;
+%unignore tensorflow::io::PyRecordReader::Close;
+%unignore tensorflow::io::PyRecordReader::New;
+
+%include "tensorflow/python/lib/io/py_record_reader.h"
+
+%unignoreall
diff --git a/tensorflow/python/lib/io/py_record_writer.cc b/tensorflow/python/lib/io/py_record_writer.cc
new file mode 100644
index 0000000000..e557756cbc
--- /dev/null
+++ b/tensorflow/python/lib/io/py_record_writer.cc
@@ -0,0 +1,44 @@
+#include "tensorflow/python/lib/io/py_record_writer.h"
+
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/platform/port.h"
+#include "tensorflow/core/lib/io/record_writer.h"
+#include "tensorflow/core/public/env.h"
+
+namespace tensorflow {
+namespace io {
+
+PyRecordWriter::PyRecordWriter() {}
+
+PyRecordWriter* PyRecordWriter::New(const string& filename) {
+ WritableFile* file;
+ Status s = Env::Default()->NewWritableFile(filename, &file);
+ if (!s.ok()) {
+ return nullptr;
+ }
+ PyRecordWriter* writer = new PyRecordWriter;
+ writer->file_ = file;
+ writer->writer_ = new RecordWriter(writer->file_);
+ return writer;
+}
+
+PyRecordWriter::~PyRecordWriter() {
+ delete writer_;
+ delete file_;
+}
+
+bool PyRecordWriter::WriteRecord(::tensorflow::StringPiece record) {
+ if (writer_ == nullptr) return false;
+ Status s = writer_->WriteRecord(record);
+ return s.ok();
+}
+
+void PyRecordWriter::Close() {
+ delete writer_;
+ delete file_;
+ writer_ = nullptr;
+ file_ = nullptr;
+}
+
+} // namespace io
+} // namespace tensorflow
diff --git a/tensorflow/python/lib/io/py_record_writer.h b/tensorflow/python/lib/io/py_record_writer.h
new file mode 100644
index 0000000000..e3fd05bd9a
--- /dev/null
+++ b/tensorflow/python/lib/io/py_record_writer.h
@@ -0,0 +1,38 @@
+#ifndef THIRD_PARTY_TENSORFLOW_PYTHON_LIB_IO_PY_RECORD_WRITER_H_
+#define THIRD_PARTY_TENSORFLOW_PYTHON_LIB_IO_PY_RECORD_WRITER_H_
+
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/platform/port.h"
+#include "tensorflow/core/public/status.h"
+
+namespace tensorflow {
+
+class WritableFile;
+
+namespace io {
+
+class RecordWriter;
+
+// A wrapper around io::RecordWriter that is more easily SWIG wrapped for
+// Python. An instance of this class is not safe for concurrent access
+// by multiple threads.
+class PyRecordWriter {
+ public:
+ static PyRecordWriter* New(const string& filename);
+ ~PyRecordWriter();
+
+ bool WriteRecord(::tensorflow::StringPiece record);
+ void Close();
+
+ private:
+ PyRecordWriter();
+
+ WritableFile* file_; // Owned
+ io::RecordWriter* writer_; // Owned
+ TF_DISALLOW_COPY_AND_ASSIGN(PyRecordWriter);
+};
+
+} // namespace io
+} // namespace tensorflow
+
+#endif // THIRD_PARTY_TENSORFLOW_PYTHON_LIB_IO_PY_RECORD_WRITER_H_
diff --git a/tensorflow/python/lib/io/py_record_writer.i b/tensorflow/python/lib/io/py_record_writer.i
new file mode 100644
index 0000000000..20fe52c495
--- /dev/null
+++ b/tensorflow/python/lib/io/py_record_writer.i
@@ -0,0 +1,38 @@
+%nothread tensorflow::io::PyRecordWriter::WriteRecord;
+
+%include "tensorflow/python/platform/base.i"
+%include "tensorflow/python/lib/core/strings.i"
+
+%feature("except") tensorflow::io::PyRecordWriter::New {
+ // Let other threads run while we write
+ Py_BEGIN_ALLOW_THREADS
+ $action
+ Py_END_ALLOW_THREADS
+}
+
+%newobject tensorflow::io::PyRecordWriter::New;
+
+%feature("except") tensorflow::io::PyRecordWriter::WriteRecord {
+ // Let other threads run while we write
+ Py_BEGIN_ALLOW_THREADS
+ $action
+ Py_END_ALLOW_THREADS
+}
+
+%{
+#include "tensorflow/python/lib/io/py_record_writer.h"
+%}
+
+%ignoreall
+
+%unignore tensorflow;
+%unignore tensorflow::io;
+%unignore tensorflow::io::PyRecordWriter;
+%unignore tensorflow::io::PyRecordWriter::~PyRecordWriter;
+%unignore tensorflow::io::PyRecordWriter::WriteRecord;
+%unignore tensorflow::io::PyRecordWriter::Close;
+%unignore tensorflow::io::PyRecordWriter::New;
+
+%include "tensorflow/python/lib/io/py_record_writer.h"
+
+%unignoreall
diff --git a/tensorflow/python/lib/io/python_io.py b/tensorflow/python/lib/io/python_io.py
new file mode 100644
index 0000000000..aedcd2ef03
--- /dev/null
+++ b/tensorflow/python/lib/io/python_io.py
@@ -0,0 +1,29 @@
+"""## Data IO (Python Functions)
+
+A TFRecords file represents a sequence of (binary) strings. The format is not
+random access, so it is suitable for streaming large amounts of data but not
+suitable if fast sharding or other non-sequential access is desired.
+
+@@TFRecordWriter
+@@tf_record_iterator
+
+- - -
+
+### TFRecords Format Details
+
+A TFRecords file contains a sequence of strings with CRC hashes. Each record
+has the format
+
+ uint64 length
+ uint32 masked_crc32_of_length
+ byte data[length]
+ uint32 masked_crc32_of_data
+
+and the records are concatenated together to produce the file. The CRC32s
+are [described here](https://en.wikipedia.org/wiki/Cyclic_redundancy_check),
+and the mask of a CRC is
+
+ masked_crc = ((crc >> 15) | (crc << 17)) + 0xa282ead8ul
+"""
+
+from tensorflow.python.lib.io.tf_record import *
diff --git a/tensorflow/python/lib/io/tf_record.py b/tensorflow/python/lib/io/tf_record.py
new file mode 100644
index 0000000000..00825bbda2
--- /dev/null
+++ b/tensorflow/python/lib/io/tf_record.py
@@ -0,0 +1,68 @@
+"""For reading and writing TFRecords files."""
+
+from tensorflow.python import pywrap_tensorflow
+
+
+def tf_record_iterator(path):
+ """An iterator that read the records from a TFRecords file.
+
+ Args:
+ path: The path to the TFRecords file.
+
+ Yields:
+ Strings.
+
+ Raises:
+ IOError: If `path` cannot be opened for reading.
+ """
+ reader = pywrap_tensorflow.PyRecordReader_New(path, 0)
+ if reader is None:
+ raise IOError("Could not open %s." % path)
+ while reader.GetNext():
+ yield reader.record()
+ reader.Close()
+
+
+class TFRecordWriter(object):
+ """A class to write records to a TFRecords file.
+
+ This class implements `__enter__` and `__exit__`, and can be used
+ in `with` blocks like a normal file.
+
+ @@__init__
+ @@write
+ @@close
+ """
+ # TODO(josh11b): Support appending?
+ def __init__(self, path):
+ """Opens file `path` and creates a `TFRecordWriter` writing to it.
+
+ Args:
+ path: The path to the TFRecords file.
+
+ Raises:
+ IOError: If `path` cannot be opened for writing.
+ """
+ self._writer = pywrap_tensorflow.PyRecordWriter_New(path)
+ if self._writer is None:
+ raise IOError("Could not write to %s." % path)
+
+ def __enter__(self):
+ """Enter a `with` block."""
+ pass
+
+ def __exit__(self, unused_type, unused_value, unused_traceback):
+ """Exit a `with` block, closing the file."""
+ self.close()
+
+ def write(self, record):
+ """Write a string record to the file.
+
+ Args:
+ record: str
+ """
+ self._writer.WriteRecord(record)
+
+ def close(self):
+ """Close the file."""
+ self._writer.Close()