chore: import upstream snapshot with attribution

tensorflow/python/lib/io/BUILD

@@ -0,0 +1,166 @@
+# python/lib/io package
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "if_oss")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test", "tf_python_pybind_extension")
+
+# copybara:uncomment_begin(google-only)
+# load("//third_party/zlib:BUILD_defs.bzl", "brittle_test_relying_on_stable_zlib_output")
+# copybara:uncomment_end
+
+visibility = [
+    "//tensorflow:__subpackages__",
+]
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = visibility,
+    licenses = ["notice"],
+)
+
+tf_python_pybind_extension(
+    name = "_pywrap_file_io",
+    srcs = ["file_io_wrapper.cc"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_pywrap_file_io.pyi",
+    ],
+    deps = [
+        "//tensorflow/core:framework_headers_lib",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_proto_parsing",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core:portable_jpeg_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:file_statistics",
+        "//tensorflow/python/lib/core:pybind11_absl",
+        "//tensorflow/python/lib/core:pybind11_status",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+        "@pybind11",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "_pywrap_record_io",
+    srcs = ["record_io_wrapper.cc"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_pywrap_record_io.pyi",
+    ],
+    deps = [
+        "//tensorflow/core:core_stringpiece",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_headers_for_pybind",
+        "//tensorflow/core:lib_proto_parsing",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:types",
+        "//tensorflow/python/lib/core:pybind11_absl",
+        "//tensorflow/python/lib/core:pybind11_status",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@pybind11",
+    ],
+)
+
+py_library(
+    name = "file_io",
+    srcs = ["file_io.py"],
+    strict_deps = True,
+    # copybara:uncomment_begin(google-only)
+    # visibility = [
+    # "//learning/serving/servables/wiz/lora:__subpackages__",
+    # "//third_party/cloud_tpu/convergence_tools:__subpackages__",
+    # "//third_party/py/tf_slim:__subpackages__",
+    # "//tensorflow:__subpackages__",
+    # "//tensorflow:internal",
+    # ],
+    # copybara:uncomment_end_and_comment_begin
+    visibility = [
+        "//visibility:public",
+    ],
+    # copybara:comment_end
+    deps = [
+        ":_pywrap_file_io",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+        "@six_archive//:six",
+    ],
+)
+
+py_library(
+    name = "tf_record",
+    srcs = ["tf_record.py"],
+    strict_deps = True,
+    visibility = [
+        "//tensorflow:__subpackages__",
+        "//tensorflow:internal",
+        "//third_party/py/tf_agents:__subpackages__",
+    ],
+    deps = [
+        ":_pywrap_record_io",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "python_io",
+    srcs = ["python_io.py"],
+    strict_deps = True,
+    visibility = [
+        "//tensorflow:__subpackages__",
+        "//tensorflow:internal",
+    ],
+    deps = [":tf_record"],
+)
+
+tf_py_strict_test(
+    name = "file_io_test",
+    size = "small",
+    srcs = ["file_io_test.py"],
+    tags = [
+        "no_rocm",
+        "no_windows",
+    ],
+    deps = [
+        ":file_io",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:gfile",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = if_oss("tf_record_test", "_tf_record_test"),
+    size = "small",
+    srcs = ["tf_record_test.py"],
+    main = "tf_record_test.py",
+    tags = [
+        # copybara:uncomment_begin(google-only)
+        # "manual",
+        # "notap",
+        # copybara:uncomment_end
+    ],
+    deps = [
+        ":tf_record",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:compat",
+    ],
+)
+
+# _tf_record_test relies on stable zlib output only during its execution.
+# It should not be brittle with respect to upgrades of zlib software or
+# different computers using different zlib software.
+# copybara:uncomment_begin(google-only)
+# brittle_test_relying_on_stable_zlib_output(
+#     name = "tf_record_test",
+#     test_target = ":_tf_record_test",
+# )
+# copybara:uncomment_end

tensorflow/python/lib/io/_pywrap_file_io.pyi

@@ -0,0 +1,53 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+class BufferedInputStream:
+    def __init__(self, filename: str, buffer_size: int) -> None: ...
+    def read(self, arg0: int) -> bytes: ...
+    def readline(self) -> bytes: ...
+    def seek(self, arg0: int) -> None: ...
+    def tell(self) -> int: ...
+
+class FileStatistics:
+    def __init__(self, *args, **kwargs) -> None: ...
+    @property
+    def is_directory(self) -> bool: ...
+    @property
+    def length(self) -> int: ...
+    @property
+    def mtime_nsec(self) -> int: ...
+
+class WritableFile:
+    def __init__(self, filename: str, mode: str) -> None: ...
+    def append(self, arg0: str) -> None: ...
+    def close(self) -> None: ...
+    def flush(self) -> None: ...
+    def tell(self) -> int: ...
+
+def CopyFile(src: str, target: str, overwrite: bool) -> None: ...
+def CreateDir(dirname: str) -> None: ...
+def DeleteFile(filename: str) -> None: ...
+def DeleteRecursively(dirname: str) -> None: ...
+def FileExists(filename: str) -> None: ...
+def GetChildren(dirname: str) -> list[str]: ...
+def GetMatchingFiles(pattern: str) -> list[str]: ...
+def GetRegisteredSchemes() -> list[str]: ...
+def HasAtomicMove(arg0: str) -> bool: ...
+def IsDirectory(dirname: str) -> bool: ...
+def ReadFileToString(filename: str) -> bytes: ...
+def RecursivelyCreateDir(dirname: str) -> None: ...
+def RenameFile(src: str, target: str, overwrite: bool) -> None: ...
+def Stat(filename: str) -> FileStatistics: ...
+def WriteStringToFile(filename: str, data: str) -> None: ...

tensorflow/python/lib/io/_pywrap_record_io.pyi

@@ -0,0 +1,52 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+class RandomRecordReader:
+    def __init__(self, arg0: str) -> None: ...
+    def close(self) -> None: ...
+    def read(self, arg0: int) -> tuple: ...
+
+class RecordIterator:
+    def __init__(self, arg0: str, arg1: str) -> None: ...
+    def close(self) -> None: ...
+    def reopen(self) -> None: ...
+    def __iter__(self) -> object: ...
+    def __next__(self) -> bytes: ...
+
+class RecordWriter:
+    def __init__(self, arg0: str, arg1: RecordWriterOptions) -> None: ...
+    def close(self) -> None: ...
+    def flush(self) -> None: ...
+    def write(self, record: str) -> None: ...
+    def __enter__(self) -> object: ...
+    def __exit__(self, *args) -> None: ...
+
+class RecordWriterOptions:
+    def __init__(self, arg0: str) -> None: ...
+    @property
+    def compression_type(self): ...
+    @property
+    def zlib_options(self) -> ZlibCompressionOptions: ...
+
+class ZlibCompressionOptions:
+    compression_level: int
+    compression_method: int
+    compression_strategy: int
+    flush_mode: int
+    input_buffer_size: int
+    mem_level: int
+    output_buffer_size: int
+    window_bits: int
+    def __init__(self, *args, **kwargs) -> None: ...

tensorflow/python/lib/io/file_io_test.py

@@ -0,0 +1,698 @@
+# This Python file uses the following encoding: utf-8
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# =============================================================================
+"""Testing File IO operations in file_io.py."""
+
+import os.path
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import errors
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import test
+
+
+class PathLike(object):
+  """Backport of pathlib.Path for Python < 3.6"""
+
+  def __init__(self, name):
+    self.name = name
+
+  def __fspath__(self):
+    return self.name
+
+  def __str__(self):
+    return self.name
+
+
+run_all_path_types = parameterized.named_parameters(
+    ("str", file_io.join),
+    ("pathlike", lambda *paths: PathLike(file_io.join(*paths))))
+
+
+class FileIoTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    self._base_dir = file_io.join(self.get_temp_dir(), "base_dir")
+    file_io.create_dir(self._base_dir)
+
+  def tearDown(self):
+    file_io.delete_recursively(self._base_dir)
+
+  def testEmptyFilename(self):
+    f = file_io.FileIO("", mode="r")
+    with self.assertRaises(errors.NotFoundError):
+      _ = f.read()
+
+  def testJoinUrlLike(self):
+    """file_io.join joins url-like filesystems with '/' on all platform."""
+    for fs in ("ram://", "gcs://", "file://"):
+      expected = fs + "exists/a/b/c.txt"
+      self.assertEqual(file_io.join(fs, "exists", "a", "b", "c.txt"), expected)
+      self.assertEqual(file_io.join(fs + "exists", "a", "b", "c.txt"), expected)
+      self.assertEqual(file_io.join(fs, "exists/a", "b", "c.txt"), expected)
+      self.assertEqual(file_io.join(fs, "exists", "a", "b/c.txt"), expected)
+
+  def testJoinFilesystem(self):
+    """file_io.join respects the os.path.join behavior for native filesystems."""
+    for sep in ("/", "\\", os.sep):
+      self.assertEqual(os.path.join("a", "b", "c"), file_io.join("a", "b", "c"))
+      self.assertEqual(
+          os.path.join(sep + "a", "b", "c"), file_io.join(sep + "a", "b", "c"))
+      self.assertEqual(
+          os.path.join("a", sep + "b", "c"), file_io.join("a", sep + "b", "c"))
+      self.assertEqual(
+          os.path.join("a", "b", sep + "c"), file_io.join("a", "b", sep + "c"))
+      self.assertEqual(
+          os.path.join("a", "b", "c" + sep), file_io.join("a", "b", "c" + sep))
+
+  @run_all_path_types
+  def testFileDoesntExist(self, join):
+    file_path = join(self._base_dir, "temp_file")
+    self.assertFalse(file_io.file_exists(file_path))
+    with self.assertRaises(errors.NotFoundError):
+      _ = file_io.read_file_to_string(file_path)
+
+  @run_all_path_types
+  def testWriteToString(self, join):
+    file_path = join(self._base_dir, "temp_file")
+    file_io.write_string_to_file(file_path, "testing")
+    self.assertTrue(file_io.file_exists(file_path))
+    file_contents = file_io.read_file_to_string(file_path)
+    self.assertEqual("testing", file_contents)
+
+  def testAtomicWriteStringToFile(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    file_io.atomic_write_string_to_file(file_path, "testing")
+    self.assertTrue(file_io.file_exists(file_path))
+    file_contents = file_io.read_file_to_string(file_path)
+    self.assertEqual("testing", file_contents)
+
+  def testAtomicWriteStringToFileOverwriteFalse(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    file_io.atomic_write_string_to_file(file_path, "old", overwrite=False)
+    with self.assertRaises(errors.AlreadyExistsError):
+      file_io.atomic_write_string_to_file(file_path, "new", overwrite=False)
+    file_contents = file_io.read_file_to_string(file_path)
+    self.assertEqual("old", file_contents)
+    file_io.delete_file(file_path)
+    file_io.atomic_write_string_to_file(file_path, "new", overwrite=False)
+    file_contents = file_io.read_file_to_string(file_path)
+    self.assertEqual("new", file_contents)
+
+  @run_all_path_types
+  def testReadBinaryMode(self, join):
+    file_path = join(self._base_dir, "temp_file")
+    file_io.write_string_to_file(file_path, "testing")
+    with file_io.FileIO(file_path, mode="rb") as f:
+      self.assertEqual(b"testing", f.read())
+
+  @run_all_path_types
+  def testWriteBinaryMode(self, join):
+    file_path = join(self._base_dir, "temp_file")
+    file_io.FileIO(file_path, "wb").write("testing")
+    with file_io.FileIO(file_path, mode="r") as f:
+      self.assertEqual("testing", f.read())
+
+  def testAppend(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with file_io.FileIO(file_path, mode="w") as f:
+      f.write("begin\n")
+    with file_io.FileIO(file_path, mode="a") as f:
+      f.write("a1\n")
+    with file_io.FileIO(file_path, mode="a") as f:
+      f.write("a2\n")
+    with file_io.FileIO(file_path, mode="r") as f:
+      file_contents = f.read()
+      self.assertEqual("begin\na1\na2\n", file_contents)
+
+  def testMultipleFiles(self):
+    file_prefix = file_io.join(self._base_dir, "temp_file")
+    for i in range(5000):
+      f = file_io.FileIO(file_prefix + str(i), mode="w+")
+      f.write("testing")
+      f.flush()
+      self.assertEqual("testing", f.read())
+      f.close()
+
+  def testMultipleWrites(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with file_io.FileIO(file_path, mode="w") as f:
+      f.write("line1\n")
+      f.write("line2")
+    file_contents = file_io.read_file_to_string(file_path)
+    self.assertEqual("line1\nline2", file_contents)
+
+  def testFileWriteBadMode(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with self.assertRaises(errors.PermissionDeniedError):
+      file_io.FileIO(file_path, mode="r").write("testing")
+
+  def testFileReadBadMode(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    file_io.FileIO(file_path, mode="w").write("testing")
+    self.assertTrue(file_io.file_exists(file_path))
+    with self.assertRaises(errors.PermissionDeniedError):
+      file_io.FileIO(file_path, mode="w").read()
+
+  @run_all_path_types
+  def testFileDelete(self, join):
+    file_path = join(self._base_dir, "temp_file")
+    file_io.FileIO(file_path, mode="w").write("testing")
+    file_io.delete_file(file_path)
+    self.assertFalse(file_io.file_exists(file_path))
+
+  def testFileDeleteFail(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with self.assertRaises(errors.NotFoundError):
+      file_io.delete_file(file_path)
+
+  def testGetMatchingFiles(self):
+    dir_path = file_io.join(self._base_dir, "temp_dir")
+    file_io.create_dir(dir_path)
+    files = ["file1.txt", "file2.txt", "file3.txt", "file*.txt"]
+    for name in files:
+      file_path = file_io.join(dir_path, name)
+      file_io.FileIO(file_path, mode="w").write("testing")
+    expected_match = [file_io.join(dir_path, name) for name in files]
+    self.assertItemsEqual(
+        file_io.get_matching_files(file_io.join(dir_path, "file*.txt")),
+        expected_match)
+    self.assertItemsEqual(file_io.get_matching_files(tuple()), [])
+    files_subset = [
+        file_io.join(dir_path, files[0]),
+        file_io.join(dir_path, files[2])
+    ]
+    self.assertItemsEqual(
+        file_io.get_matching_files(files_subset), files_subset)
+    file_io.delete_recursively(dir_path)
+    self.assertFalse(file_io.file_exists(file_io.join(dir_path, "file3.txt")))
+
+  def testGetMatchingFilesWhenParentDirContainsParantheses(self):
+    dir_path = file_io.join(self._base_dir, "dir_(special)")
+    file_io.create_dir(dir_path)
+    files = ["file1.txt", "file(2).txt"]
+    for name in files:
+      file_path = file_io.join(dir_path, name)
+      file_io.FileIO(file_path, mode="w").write("testing")
+    expected_match = [file_io.join(dir_path, name) for name in files]
+    glob_pattern = file_io.join(dir_path, "*")
+    self.assertItemsEqual(
+        file_io.get_matching_files(glob_pattern), expected_match)
+
+  @run_all_path_types
+  def testCreateRecursiveDir(self, join):
+    dir_path = join(self._base_dir, "temp_dir/temp_dir1/temp_dir2")
+    file_io.recursive_create_dir(dir_path)
+    file_io.recursive_create_dir(dir_path)  # repeat creation
+    file_path = file_io.join(str(dir_path), "temp_file")
+    file_io.FileIO(file_path, mode="w").write("testing")
+    self.assertTrue(file_io.file_exists(file_path))
+    file_io.delete_recursively(file_io.join(self._base_dir, "temp_dir"))
+    self.assertFalse(file_io.file_exists(file_path))
+
+  @run_all_path_types
+  def testCopy(self, join):
+    file_path = join(self._base_dir, "temp_file")
+    file_io.FileIO(file_path, mode="w").write("testing")
+    copy_path = join(self._base_dir, "copy_file")
+    file_io.copy(file_path, copy_path)
+    self.assertTrue(file_io.file_exists(copy_path))
+    f = file_io.FileIO(file_path, mode="r")
+    self.assertEqual("testing", f.read())
+    self.assertEqual(7, f.tell())
+
+  def testCopyOverwrite(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    file_io.FileIO(file_path, mode="w").write("testing")
+    copy_path = file_io.join(self._base_dir, "copy_file")
+    file_io.FileIO(copy_path, mode="w").write("copy")
+    file_io.copy(file_path, copy_path, overwrite=True)
+    self.assertTrue(file_io.file_exists(copy_path))
+    self.assertEqual("testing", file_io.FileIO(file_path, mode="r").read())
+
+  def testCopyOverwriteFalse(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    file_io.FileIO(file_path, mode="w").write("testing")
+    copy_path = file_io.join(self._base_dir, "copy_file")
+    file_io.FileIO(copy_path, mode="w").write("copy")
+    with self.assertRaises(errors.AlreadyExistsError):
+      file_io.copy(file_path, copy_path, overwrite=False)
+
+  @run_all_path_types
+  def testRename(self, join):
+    file_path = join(self._base_dir, "temp_file")
+    file_io.FileIO(file_path, mode="w").write("testing")
+    rename_path = join(self._base_dir, "rename_file")
+    file_io.rename(file_path, rename_path)
+    self.assertTrue(file_io.file_exists(rename_path))
+    self.assertFalse(file_io.file_exists(file_path))
+
+  def testRenameOverwrite(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    file_io.FileIO(file_path, mode="w").write("testing")
+    rename_path = file_io.join(self._base_dir, "rename_file")
+    file_io.FileIO(rename_path, mode="w").write("rename")
+    file_io.rename(file_path, rename_path, overwrite=True)
+    self.assertTrue(file_io.file_exists(rename_path))
+    self.assertFalse(file_io.file_exists(file_path))
+
+  def testRenameOverwriteFalse(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    file_io.FileIO(file_path, mode="w").write("testing")
+    rename_path = file_io.join(self._base_dir, "rename_file")
+    file_io.FileIO(rename_path, mode="w").write("rename")
+    with self.assertRaises(errors.AlreadyExistsError):
+      file_io.rename(file_path, rename_path, overwrite=False)
+    self.assertTrue(file_io.file_exists(rename_path))
+    self.assertTrue(file_io.file_exists(file_path))
+
+  def testDeleteRecursivelyFail(self):
+    fake_dir_path = file_io.join(self._base_dir, "temp_dir")
+    with self.assertRaises(errors.NotFoundError):
+      file_io.delete_recursively(fake_dir_path)
+
+  @run_all_path_types
+  def testIsDirectory(self, join):
+    dir_path = join(self._base_dir, "test_dir")
+    # Failure for a non-existing dir.
+    self.assertFalse(file_io.is_directory(dir_path))
+    file_io.create_dir(dir_path)
+    self.assertTrue(file_io.is_directory(dir_path))
+    file_path = join(str(dir_path), "test_file")
+    file_io.FileIO(file_path, mode="w").write("test")
+    # False for a file.
+    self.assertFalse(file_io.is_directory(file_path))
+    # Test that the value returned from `stat()` has `is_directory` set.
+    file_statistics = file_io.stat(dir_path)
+    self.assertTrue(file_statistics.is_directory)
+
+  @run_all_path_types
+  def testListDirectory(self, join):
+    dir_path = join(self._base_dir, "test_dir")
+    file_io.create_dir(dir_path)
+    files = ["file1.txt", "file2.txt", "file3.txt"]
+    for name in files:
+      file_path = join(str(dir_path), name)
+      file_io.FileIO(file_path, mode="w").write("testing")
+    subdir_path = join(str(dir_path), "sub_dir")
+    file_io.create_dir(subdir_path)
+    subdir_file_path = join(str(subdir_path), "file4.txt")
+    file_io.FileIO(subdir_file_path, mode="w").write("testing")
+    dir_list = file_io.list_directory(dir_path)
+    self.assertItemsEqual(files + ["sub_dir"], dir_list)
+
+  def testListDirectoryFailure(self):
+    dir_path = file_io.join(self._base_dir, "test_dir")
+    with self.assertRaises(errors.NotFoundError):
+      file_io.list_directory(dir_path)
+
+  def _setupWalkDirectories(self, dir_path):
+    # Creating a file structure as follows
+    # test_dir -> file: file1.txt; dirs: subdir1_1, subdir1_2, subdir1_3
+    # subdir1_1 -> file: file3.txt
+    # subdir1_2 -> dir: subdir2
+    file_io.create_dir(dir_path)
+    file_io.FileIO(
+        file_io.join(dir_path, "file1.txt"), mode="w").write("testing")
+    sub_dirs1 = ["subdir1_1", "subdir1_2", "subdir1_3"]
+    for name in sub_dirs1:
+      file_io.create_dir(file_io.join(dir_path, name))
+    file_io.FileIO(
+        file_io.join(dir_path, "subdir1_1/file2.txt"),
+        mode="w").write("testing")
+    file_io.create_dir(file_io.join(dir_path, "subdir1_2/subdir2"))
+
+  @run_all_path_types
+  def testWalkInOrder(self, join):
+    dir_path_str = file_io.join(self._base_dir, "test_dir")
+    dir_path = join(self._base_dir, "test_dir")
+    self._setupWalkDirectories(dir_path_str)
+    # Now test the walk (in_order = True)
+    all_dirs = []
+    all_subdirs = []
+    all_files = []
+    for (w_dir, w_subdirs, w_files) in file_io.walk(dir_path, in_order=True):
+      all_dirs.append(w_dir)
+      all_subdirs.append(w_subdirs)
+      all_files.append(w_files)
+    self.assertItemsEqual(all_dirs, [dir_path_str] + [
+        file_io.join(dir_path_str, item) for item in
+        ["subdir1_1", "subdir1_2", "subdir1_2/subdir2", "subdir1_3"]
+    ])
+    self.assertEqual(dir_path_str, all_dirs[0])
+    self.assertLess(
+        all_dirs.index(file_io.join(dir_path_str, "subdir1_2")),
+        all_dirs.index(file_io.join(dir_path_str, "subdir1_2/subdir2")))
+    self.assertItemsEqual(all_subdirs[1:5], [[], ["subdir2"], [], []])
+    self.assertItemsEqual(all_subdirs[0],
+                          ["subdir1_1", "subdir1_2", "subdir1_3"])
+    self.assertItemsEqual(all_files, [["file1.txt"], ["file2.txt"], [], [], []])
+    self.assertLess(
+        all_files.index(["file1.txt"]), all_files.index(["file2.txt"]))
+
+  def testWalkPostOrder(self):
+    dir_path = file_io.join(self._base_dir, "test_dir")
+    self._setupWalkDirectories(dir_path)
+    # Now test the walk (in_order = False)
+    all_dirs = []
+    all_subdirs = []
+    all_files = []
+    for (w_dir, w_subdirs, w_files) in file_io.walk(dir_path, in_order=False):
+      all_dirs.append(w_dir)
+      all_subdirs.append(w_subdirs)
+      all_files.append(w_files)
+    self.assertItemsEqual(all_dirs, [
+        file_io.join(dir_path, item) for item in
+        ["subdir1_1", "subdir1_2/subdir2", "subdir1_2", "subdir1_3"]
+    ] + [dir_path])
+    self.assertEqual(dir_path, all_dirs[4])
+    self.assertLess(
+        all_dirs.index(file_io.join(dir_path, "subdir1_2/subdir2")),
+        all_dirs.index(file_io.join(dir_path, "subdir1_2")))
+    self.assertItemsEqual(all_subdirs[0:4], [[], [], ["subdir2"], []])
+    self.assertItemsEqual(all_subdirs[4],
+                          ["subdir1_1", "subdir1_2", "subdir1_3"])
+    self.assertItemsEqual(all_files, [["file2.txt"], [], [], [], ["file1.txt"]])
+    self.assertLess(
+        all_files.index(["file2.txt"]), all_files.index(["file1.txt"]))
+
+  def testWalkFailure(self):
+    dir_path = file_io.join(self._base_dir, "test_dir")
+    # Try walking a directory that wasn't created.
+    all_dirs = []
+    all_subdirs = []
+    all_files = []
+    for (w_dir, w_subdirs, w_files) in file_io.walk(dir_path, in_order=False):
+      all_dirs.append(w_dir)
+      all_subdirs.append(w_subdirs)
+      all_files.append(w_files)
+    self.assertItemsEqual(all_dirs, [])
+    self.assertItemsEqual(all_subdirs, [])
+    self.assertItemsEqual(all_files, [])
+
+  @run_all_path_types
+  def testStat(self, join):
+    file_path = join(self._base_dir, "temp_file")
+    file_io.FileIO(file_path, mode="w").write("testing")
+    file_statistics = file_io.stat(file_path)
+    os_statistics = os.stat(str(file_path))
+    self.assertEqual(7, file_statistics.length)
+    self.assertEqual(
+        int(os_statistics.st_mtime), int(file_statistics.mtime_nsec / 1e9))
+    self.assertFalse(file_statistics.is_directory)
+
+  def testReadLine(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with file_io.FileIO(file_path, mode="r+") as f:
+      f.write("testing1\ntesting2\ntesting3\n\ntesting5")
+    self.assertEqual(36, f.size())
+    self.assertEqual("testing1\n", f.readline())
+    self.assertEqual("testing2\n", f.readline())
+    self.assertEqual("testing3\n", f.readline())
+    self.assertEqual("\n", f.readline())
+    self.assertEqual("testing5", f.readline())
+    self.assertEqual("", f.readline())
+
+  def testRead(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with file_io.FileIO(file_path, mode="r+") as f:
+      f.write("testing1\ntesting2\ntesting3\n\ntesting5")
+    self.assertEqual(36, f.size())
+    self.assertEqual("testing1\n", f.read(9))
+    self.assertEqual("testing2\n", f.read(9))
+    self.assertEqual("t", f.read(1))
+    self.assertEqual("esting3\n\ntesting5", f.read())
+
+  def testReadErrorReacquiresGil(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with file_io.FileIO(file_path, mode="r+") as f:
+      f.write("testing1\ntesting2\ntesting3\n\ntesting5")
+    with self.assertRaises(errors.InvalidArgumentError):
+      # At present, this is sufficient to convince ourselves that the change
+      # fixes the problem. That is, this test will seg fault without the change,
+      # and pass with it. Unfortunately, this is brittle, as it relies on the
+      # Python layer to pass the argument along to the wrapped C++ without
+      # checking the argument itself.
+      f.read(-2)
+
+  def testTell(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with file_io.FileIO(file_path, mode="r+") as f:
+      f.write("testing1\ntesting2\ntesting3\n\ntesting5")
+    self.assertEqual(0, f.tell())
+    self.assertEqual("testing1\n", f.readline())
+    self.assertEqual(9, f.tell())
+    self.assertEqual("testing2\n", f.readline())
+    self.assertEqual(18, f.tell())
+    self.assertEqual("testing3\n", f.readline())
+    self.assertEqual(27, f.tell())
+    self.assertEqual("\n", f.readline())
+    self.assertEqual(28, f.tell())
+    self.assertEqual("testing5", f.readline())
+    self.assertEqual(36, f.tell())
+    self.assertEqual("", f.readline())
+    self.assertEqual(36, f.tell())
+
+  def testSeek(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with file_io.FileIO(file_path, mode="r+") as f:
+      f.write("testing1\ntesting2\ntesting3\n\ntesting5")
+    self.assertEqual("testing1\n", f.readline())
+    self.assertEqual(9, f.tell())
+
+    # Seek to 18
+    f.seek(18)
+    self.assertEqual(18, f.tell())
+    self.assertEqual("testing3\n", f.readline())
+
+    # Seek back to 9
+    f.seek(9)
+    self.assertEqual(9, f.tell())
+    self.assertEqual("testing2\n", f.readline())
+
+    f.seek(0)
+    self.assertEqual(0, f.tell())
+    self.assertEqual("testing1\n", f.readline())
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      f.seek(-1)
+
+    with self.assertRaises(TypeError):
+      f.seek()
+
+    # TODO(jhseu): Delete after position deprecation.
+    with self.assertRaises(TypeError):
+      f.seek(offset=0, position=0)
+    f.seek(position=9)
+    self.assertEqual(9, f.tell())
+    self.assertEqual("testing2\n", f.readline())
+
+  def testSeekFromWhat(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with file_io.FileIO(file_path, mode="r+") as f:
+      f.write("testing1\ntesting2\ntesting3\n\ntesting5")
+    self.assertEqual("testing1\n", f.readline())
+    self.assertEqual(9, f.tell())
+
+    # Seek to 18
+    f.seek(9, 1)
+    self.assertEqual(18, f.tell())
+    self.assertEqual("testing3\n", f.readline())
+
+    # Seek back to 9
+    f.seek(9, 0)
+    self.assertEqual(9, f.tell())
+    self.assertEqual("testing2\n", f.readline())
+
+    f.seek(-f.size(), 2)
+    self.assertEqual(0, f.tell())
+    self.assertEqual("testing1\n", f.readline())
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      f.seek(0, 3)
+
+  def testReadingIterator(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    data = ["testing1\n", "testing2\n", "testing3\n", "\n", "testing5"]
+    with file_io.FileIO(file_path, mode="r+") as f:
+      f.write("".join(data))
+    actual_data = []
+    for line in f:
+      actual_data.append(line)
+    self.assertSequenceEqual(actual_data, data)
+
+  def testReadlines(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    data = ["testing1\n", "testing2\n", "testing3\n", "\n", "testing5"]
+    f = file_io.FileIO(file_path, mode="r+")
+    f.write("".join(data))
+    f.flush()
+    lines = f.readlines()
+    self.assertSequenceEqual(lines, data)
+
+  def testUTF8StringPath(self):
+    file_path = file_io.join(self._base_dir, "UTF8测试_file")
+    file_io.write_string_to_file(file_path, "testing")
+    with file_io.FileIO(file_path, mode="rb") as f:
+      self.assertEqual(b"testing", f.read())
+
+  def testEof(self):
+    """Test that reading past EOF does not raise an exception."""
+
+    file_path = file_io.join(self._base_dir, "temp_file")
+    f = file_io.FileIO(file_path, mode="r+")
+    content = "testing"
+    f.write(content)
+    f.flush()
+    self.assertEqual(content, f.read(len(content) + 1))
+
+  @run_all_path_types
+  def testUTF8StringPathExists(self, join):
+    file_path = join(self._base_dir, "UTF8测试_file_exist")
+    file_io.write_string_to_file(file_path, "testing")
+    v = file_io.file_exists(file_path)
+    self.assertEqual(v, True)
+
+  def testFilecmp(self):
+    file1 = file_io.join(self._base_dir, "file1")
+    file_io.write_string_to_file(file1, "This is a sentence\n" * 100)
+
+    file2 = file_io.join(self._base_dir, "file2")
+    file_io.write_string_to_file(file2, "This is another sentence\n" * 100)
+
+    file3 = file_io.join(self._base_dir, "file3")
+    file_io.write_string_to_file(file3, u"This is another sentence\n" * 100)
+
+    self.assertFalse(file_io.filecmp(file1, file2))
+    self.assertTrue(file_io.filecmp(file2, file3))
+
+  def testFilecmpSameSize(self):
+    file1 = file_io.join(self._base_dir, "file1")
+    file_io.write_string_to_file(file1, "This is a sentence\n" * 100)
+
+    file2 = file_io.join(self._base_dir, "file2")
+    file_io.write_string_to_file(file2, "This is b sentence\n" * 100)
+
+    file3 = file_io.join(self._base_dir, "file3")
+    file_io.write_string_to_file(file3, u"This is b sentence\n" * 100)
+
+    self.assertFalse(file_io.filecmp(file1, file2))
+    self.assertTrue(file_io.filecmp(file2, file3))
+
+  def testFilecmpBinary(self):
+    file1 = file_io.join(self._base_dir, "file1")
+    file_io.FileIO(file1, "wb").write("testing\n\na")
+
+    file2 = file_io.join(self._base_dir, "file2")
+    file_io.FileIO(file2, "wb").write("testing\n\nb")
+
+    file3 = file_io.join(self._base_dir, "file3")
+    file_io.FileIO(file3, "wb").write("testing\n\nb")
+
+    file4 = file_io.join(self._base_dir, "file4")
+    file_io.FileIO(file4, "wb").write("testing\n\ntesting")
+
+    self.assertFalse(file_io.filecmp(file1, file2))
+    self.assertFalse(file_io.filecmp(file1, file4))
+    self.assertTrue(file_io.filecmp(file2, file3))
+
+  def testFileCrc32(self):
+    file1 = file_io.join(self._base_dir, "file1")
+    file_io.write_string_to_file(file1, "This is a sentence\n" * 100)
+    crc1 = file_io.file_crc32(file1)
+
+    file2 = file_io.join(self._base_dir, "file2")
+    file_io.write_string_to_file(file2, "This is another sentence\n" * 100)
+    crc2 = file_io.file_crc32(file2)
+
+    file3 = file_io.join(self._base_dir, "file3")
+    file_io.write_string_to_file(file3, "This is another sentence\n" * 100)
+    crc3 = file_io.file_crc32(file3)
+
+    self.assertTrue(crc1 != crc2)
+    self.assertEqual(crc2, crc3)
+
+  def testFileCrc32WithBytes(self):
+    file1 = file_io.join(self._base_dir, "file1")
+    file_io.write_string_to_file(file1, "This is a sentence\n" * 100)
+    crc1 = file_io.file_crc32(file1, block_size=24)
+
+    file2 = file_io.join(self._base_dir, "file2")
+    file_io.write_string_to_file(file2, "This is another sentence\n" * 100)
+    crc2 = file_io.file_crc32(file2, block_size=24)
+
+    file3 = file_io.join(self._base_dir, "file3")
+    file_io.write_string_to_file(file3, "This is another sentence\n" * 100)
+    crc3 = file_io.file_crc32(file3, block_size=-1)
+
+    self.assertTrue(crc1 != crc2)
+    self.assertEqual(crc2, crc3)
+
+  def testFileCrc32Binary(self):
+    file1 = file_io.join(self._base_dir, "file1")
+    file_io.FileIO(file1, "wb").write("testing\n\n")
+    crc1 = file_io.file_crc32(file1)
+
+    file2 = file_io.join(self._base_dir, "file2")
+    file_io.FileIO(file2, "wb").write("testing\n\n\n")
+    crc2 = file_io.file_crc32(file2)
+
+    file3 = file_io.join(self._base_dir, "file3")
+    file_io.FileIO(file3, "wb").write("testing\n\n\n")
+    crc3 = file_io.file_crc32(file3)
+
+    self.assertTrue(crc1 != crc2)
+    self.assertEqual(crc2, crc3)
+
+  def testFileSeekableWithZip(self):
+    # Note: Test case for GitHub issue 27276, issue only exposed in python 3.7+.
+    filename = file_io.join(self._base_dir, "a.npz")
+    np.savez_compressed(filename, {"a": 1, "b": 2})
+    with gfile.GFile(filename, "rb") as f:
+      info = np.load(f, allow_pickle=True)  # pylint: disable=unexpected-keyword-arg
+    _ = [i for i in info.items()]
+
+  def testHasAtomicMove(self):
+    self.assertTrue(file_io.has_atomic_move("/a/b/c"))
+
+  def testGetRegisteredSchemes(self):
+    expected = ["", "file", "ram"]
+    actual = file_io.get_registered_schemes()
+    # Be flexible about additional schemes that may sometimes be registered when
+    # this test is run, while still verifying each scheme appears just once.
+    maybe_expected = ["gs", "hypercomputer"]
+    for scheme in maybe_expected:
+      if scheme in actual:
+        expected.append(scheme)
+    self.assertCountEqual(expected, actual)
+
+  def testReadWriteWithEncoding(self):
+    file_path = file_io.join(self._base_dir, "temp_file")
+    with open(file_path, mode="w", encoding="cp932") as f:
+      f.write("今日はいい天気")
+    with file_io.FileIO(file_path, mode="r", encoding="cp932") as f:
+      self.assertEqual(f.read(), "今日はいい天気")
+
+    with file_io.FileIO(file_path, mode="w", encoding="cp932") as f:
+      f.write("今日はいい天気")
+    with open(file_path, mode="r", encoding="cp932") as f:
+      self.assertEqual(f.read(), "今日はいい天気")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/lib/io/file_io_wrapper.cc

@@ -0,0 +1,310 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/core/lib/core/error_codes.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/io/buffered_inputstream.h"
+#include "tensorflow/core/lib/io/random_inputstream.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/file_statistics.h"
+#include "tensorflow/core/platform/file_system.h"
+#include "tensorflow/core/platform/stringpiece.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/python/lib/core/pybind11_absl.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+namespace tensorflow {
+}  // namespace tensorflow
+
+namespace {
+namespace py = pybind11;
+
+PYBIND11_MODULE(_pywrap_file_io, m) {
+  m.def(
+      "FileExists",
+      [](const std::string& filename) {
+        absl::Status status;
+        {
+          py::gil_scoped_release release;
+          status = tensorflow::Env::Default()->FileExists(filename);
+        }
+        tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+      },
+      py::arg("filename"));
+  m.def(
+      "DeleteFile",
+      [](const std::string& filename) {
+        py::gil_scoped_release release;
+        absl::Status status = tensorflow::Env::Default()->DeleteFile(filename);
+        tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+      },
+      py::arg("filename"));
+  m.def(
+      "ReadFileToString",
+      [](const std::string& filename) {
+        std::string data;
+        py::gil_scoped_release release;
+        const auto status =
+            ReadFileToString(tensorflow::Env::Default(), filename, &data);
+        pybind11::gil_scoped_acquire acquire;
+        tensorflow::MaybeRaiseRegisteredFromStatus(status);
+        return py::bytes(data);
+      },
+      py::arg("filename"));
+  m.def(
+      "WriteStringToFile",
+      [](const std::string& filename, absl::string_view data) {
+        py::gil_scoped_release release;
+        const auto status =
+            WriteStringToFile(tensorflow::Env::Default(), filename, data);
+        tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+      },
+      py::arg("filename"), py::arg("data"));
+  m.def(
+      "GetChildren",
+      [](const std::string& dirname) {
+        std::vector<std::string> results;
+        py::gil_scoped_release release;
+        const auto status =
+            tensorflow::Env::Default()->GetChildren(dirname, &results);
+        pybind11::gil_scoped_acquire acquire;
+        tensorflow::MaybeRaiseRegisteredFromStatus(status);
+        return results;
+      },
+      py::arg("dirname"));
+  m.def(
+      "GetMatchingFiles",
+      [](const std::string& pattern) {
+        std::vector<std::string> results;
+        py::gil_scoped_release release;
+        const auto status =
+            tensorflow::Env::Default()->GetMatchingPaths(pattern, &results);
+        pybind11::gil_scoped_acquire acquire;
+        tensorflow::MaybeRaiseRegisteredFromStatus(status);
+        return results;
+      },
+      py::arg("pattern"));
+  m.def(
+      "CreateDir",
+      [](const std::string& dirname) {
+        py::gil_scoped_release release;
+        const auto status = tensorflow::Env::Default()->CreateDir(dirname);
+        if (absl::IsAlreadyExists(status)) {
+          return;
+        }
+        tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+      },
+      py::arg("dirname"));
+  m.def(
+      "RecursivelyCreateDir",
+      [](const std::string& dirname) {
+        py::gil_scoped_release release;
+        const auto status =
+            tensorflow::Env::Default()->RecursivelyCreateDir(dirname);
+        tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+      },
+      py::arg("dirname"));
+  m.def(
+      "CopyFile",
+      [](const std::string& src, const std::string& target, bool overwrite) {
+        py::gil_scoped_release release;
+        auto* env = tensorflow::Env::Default();
+        absl::Status status;
+        if (!overwrite && env->FileExists(target).ok()) {
+          status = absl::AlreadyExistsError("file already exists");
+        } else {
+          status = env->CopyFile(src, target);
+        }
+        tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+      },
+      py::arg("src"), py::arg("target"), py::arg("overwrite"));
+  m.def(
+      "RenameFile",
+      [](const std::string& src, const std::string& target, bool overwrite) {
+        py::gil_scoped_release release;
+        auto* env = tensorflow::Env::Default();
+        absl::Status status;
+        if (!overwrite && env->FileExists(target).ok()) {
+          status = absl::AlreadyExistsError("file already exists");
+        } else {
+          status = env->RenameFile(src, target);
+        }
+        tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+      },
+      py::arg("src"), py::arg("target"), py::arg("overwrite"));
+  m.def(
+      "DeleteRecursively",
+      [](const std::string& dirname) {
+        py::gil_scoped_release release;
+        int64_t undeleted_files;
+        int64_t undeleted_dirs;
+        auto status = tensorflow::Env::Default()->DeleteRecursively(
+            dirname, &undeleted_files, &undeleted_dirs);
+        if (status.ok() && (undeleted_files > 0 || undeleted_dirs > 0)) {
+          status = absl::PermissionDeniedError("could not fully delete dir");
+        }
+        tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+      },
+      py::arg("dirname"));
+  m.def(
+      "IsDirectory",
+      [](const std::string& dirname) {
+        py::gil_scoped_release release;
+        const auto status = tensorflow::Env::Default()->IsDirectory(dirname);
+        // FAILED_PRECONDITION response means path exists but isn't a dir.
+        if (absl::IsFailedPrecondition(status)) {
+          return false;
+        }
+
+        tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+        return true;
+      },
+      py::arg("dirname"));
+  m.def("HasAtomicMove", [](const std::string& path) {
+    py::gil_scoped_release release;
+    bool has_atomic_move;
+    const auto status =
+        tensorflow::Env::Default()->HasAtomicMove(path, &has_atomic_move);
+    tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+    return has_atomic_move;
+  });
+
+  py::class_<tensorflow::FileStatistics>(m, "FileStatistics")
+      .def_readonly("length", &tensorflow::FileStatistics::length)
+      .def_readonly("mtime_nsec", &tensorflow::FileStatistics::mtime_nsec)
+      .def_readonly("is_directory", &tensorflow::FileStatistics::is_directory);
+
+  m.def(
+      "Stat",
+      [](const std::string& filename) {
+        py::gil_scoped_release release;
+        std::unique_ptr<tensorflow::FileStatistics> self(
+            new tensorflow::FileStatistics);
+        const auto status =
+            tensorflow::Env::Default()->Stat(filename, self.get());
+        py::gil_scoped_acquire acquire;
+        tensorflow::MaybeRaiseRegisteredFromStatus(status);
+        return self.release();
+      },
+      py::arg("filename"));
+
+  m.def("GetRegisteredSchemes", []() {
+    std::vector<std::string> results;
+    py::gil_scoped_release release;
+    const auto status =
+        tensorflow::Env::Default()->GetRegisteredFileSystemSchemes(&results);
+    pybind11::gil_scoped_acquire acquire;
+    tensorflow::MaybeRaiseRegisteredFromStatus(status);
+    return results;
+  });
+
+  using tensorflow::WritableFile;
+  py::class_<WritableFile>(m, "WritableFile")
+      .def(py::init([](const std::string& filename, const std::string& mode) {
+             py::gil_scoped_release release;
+             auto* env = tensorflow::Env::Default();
+             std::unique_ptr<WritableFile> self;
+             const auto status = mode.find('a') == std::string::npos
+                                     ? env->NewWritableFile(filename, &self)
+                                     : env->NewAppendableFile(filename, &self);
+             py::gil_scoped_acquire acquire;
+             tensorflow::MaybeRaiseRegisteredFromStatus(status);
+             return self.release();
+           }),
+           py::arg("filename"), py::arg("mode"))
+      .def("append",
+           [](WritableFile* self, absl::string_view data) {
+             const auto status = self->Append(data);
+             tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+           })
+      // TODO(slebedev): Make WritableFile::Tell const and change self
+      // to be a reference.
+      .def("tell",
+           [](WritableFile* self) {
+             int64_t pos = -1;
+             py::gil_scoped_release release;
+             const auto status = self->Tell(&pos);
+             tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+             return pos;
+           })
+      .def("flush",
+           [](WritableFile* self) {
+             py::gil_scoped_release release;
+             tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(self->Flush());
+           })
+      .def("close", [](WritableFile* self) {
+        py::gil_scoped_release release;
+        tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(self->Close());
+      });
+
+  using tensorflow::io::BufferedInputStream;
+  py::class_<BufferedInputStream>(m, "BufferedInputStream")
+      .def(py::init([](const std::string& filename, size_t buffer_size) {
+             py::gil_scoped_release release;
+             std::unique_ptr<tensorflow::RandomAccessFile> file;
+             const auto status =
+                 tensorflow::Env::Default()->NewRandomAccessFile(filename,
+                                                                 &file);
+             tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+             std::unique_ptr<tensorflow::io::RandomAccessInputStream>
+                 input_stream(new tensorflow::io::RandomAccessInputStream(
+                     file.release(),
+                     /*owns_file=*/true));
+             py::gil_scoped_acquire acquire;
+             return new BufferedInputStream(input_stream.release(), buffer_size,
+                                            /*owns_input_stream=*/true);
+           }),
+           py::arg("filename"), py::arg("buffer_size"))
+      .def("read",
+           [](BufferedInputStream* self, int64_t bytes_to_read) {
+             py::gil_scoped_release release;
+             tensorflow::tstring result;
+             const auto status = self->ReadNBytes(bytes_to_read, &result);
+             if (!status.ok() && !absl::IsOutOfRange(status)) {
+               result.clear();
+               tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
+             }
+             py::gil_scoped_acquire acquire;
+             return py::bytes(result);
+           })
+      .def("readline",
+           [](BufferedInputStream* self) {
+             py::gil_scoped_release release;
+             auto output = self->ReadLineAsString();
+             py::gil_scoped_acquire acquire;
+             return py::bytes(output);
+           })
+      .def("seek",
+           [](BufferedInputStream* self, int64_t pos) {
+             py::gil_scoped_release release;
+             tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(self->Seek(pos));
+           })
+      .def("tell", [](BufferedInputStream* self) {
+        py::gil_scoped_release release;
+        return self->Tell();
+      });
+}
+}  // namespace

tensorflow/python/lib/io/python_io.py

@@ -0,0 +1,24 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+"""Python functions for directly manipulating TFRecord-formatted files.
+
+API docstring: tensorflow.python_io
+"""
+
+# go/tf-wildcard-import
+# pylint: disable=wildcard-import
+from tensorflow.python.lib.io.tf_record import *
+# pylint: enable=wildcard-import

tensorflow/python/lib/io/record_io_wrapper.cc

@@ -0,0 +1,426 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/io/record_reader.h"
+#include "tensorflow/core/lib/io/record_writer.h"
+#include "tensorflow/core/lib/io/zlib_compression_options.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/file_system.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/python/lib/core/pybind11_absl.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+namespace {
+
+namespace py = ::pybind11;
+
+class PyRecordReader {
+ public:
+  // NOTE(sethtroisi): At this time PyRecordReader doesn't benefit from taking
+  // RecordReaderOptions, if this changes the API can be updated at that time.
+  static absl::Status New(const std::string& filename,
+                          const std::string& compression_type,
+                          PyRecordReader** out) {
+    auto tmp = new PyRecordReader(filename, compression_type);
+    TF_RETURN_IF_ERROR(tmp->Reopen());
+    *out = tmp;
+    return absl::OkStatus();
+  }
+
+  PyRecordReader() = delete;
+  ~PyRecordReader() { Close(); }
+
+  absl::Status ReadNextRecord(tensorflow::tstring* out) {
+    tensorflow::mutex_lock lock(mutex_);
+    if (IsClosedInternal()) {
+      return absl::FailedPreconditionError("Reader is closed.");
+    }
+    return reader_->ReadRecord(&offset_, out);
+  }
+
+  bool IsClosed() const {
+    tensorflow::mutex_lock lock(mutex_);
+    return IsClosedInternal();
+  }
+
+  void Close() {
+    tensorflow::mutex_lock lock(mutex_);
+    reader_ = nullptr;
+    file_ = nullptr;
+  }
+
+  // Reopen a closed writer by re-opening the file and re-creating the reader,
+  // but preserving the prior read offset. If not closed, returns an error.
+  //
+  // This is useful to allow "refreshing" the underlying file handle, in cases
+  // where the file was replaced with a newer version containing additional data
+  // that otherwise wouldn't be available via the existing file handle. This
+  // allows the file to be polled continuously using the same iterator, even as
+  // it grows, which supports use cases such as TensorBoard.
+  absl::Status Reopen() {
+    tensorflow::mutex_lock lock(mutex_);
+    if (!IsClosedInternal()) {
+      return absl::FailedPreconditionError("Reader is not closed.");
+    }
+    TF_RETURN_IF_ERROR(
+        tensorflow::Env::Default()->NewRandomAccessFile(filename_, &file_));
+    reader_ =
+        std::make_unique<tensorflow::io::RecordReader>(file_.get(), options_);
+    return absl::OkStatus();
+  }
+
+ private:
+  bool IsClosedInternal() const {
+    return file_ == nullptr && reader_ == nullptr;
+  }
+
+  static constexpr uint64_t kReaderBufferSize = 16 * 1024 * 1024;
+
+  PyRecordReader(const std::string& filename,
+                 const std::string& compression_type)
+      : filename_(filename),
+        options_(CreateOptions(compression_type)),
+        offset_(0),
+        file_(nullptr),
+        reader_(nullptr) {}
+
+  static tensorflow::io::RecordReaderOptions CreateOptions(
+      const std::string& compression_type) {
+    auto options =
+        tensorflow::io::RecordReaderOptions::CreateRecordReaderOptions(
+            compression_type);
+    options.buffer_size = kReaderBufferSize;
+    return options;
+  }
+
+  const std::string filename_;
+  const tensorflow::io::RecordReaderOptions options_;
+  uint64_t offset_;
+  std::unique_ptr<tensorflow::RandomAccessFile> file_;
+  std::unique_ptr<tensorflow::io::RecordReader> reader_;
+  mutable tensorflow::mutex mutex_;
+
+  PyRecordReader(const PyRecordReader&) = delete;
+  void operator=(const PyRecordReader&) = delete;
+};
+
+class PyRecordRandomReader {
+ public:
+  static absl::Status New(const std::string& filename,
+                          PyRecordRandomReader** out) {
+    std::unique_ptr<tensorflow::RandomAccessFile> file;
+    TF_RETURN_IF_ERROR(
+        tensorflow::Env::Default()->NewRandomAccessFile(filename, &file));
+    auto options =
+        tensorflow::io::RecordReaderOptions::CreateRecordReaderOptions("");
+    options.buffer_size = kReaderBufferSize;
+    auto reader =
+        std::make_unique<tensorflow::io::RecordReader>(file.get(), options);
+    *out = new PyRecordRandomReader(std::move(file), std::move(reader));
+    return absl::OkStatus();
+  }
+
+  PyRecordRandomReader() = delete;
+  ~PyRecordRandomReader() { Close(); }
+
+  absl::Status ReadRecord(uint64_t* offset, tensorflow::tstring* out) {
+    tensorflow::mutex_lock lock(mutex_);
+    if (IsClosedInternal()) {
+      return absl::FailedPreconditionError("Random TFRecord Reader is closed.");
+    }
+    return reader_->ReadRecord(offset, out);
+  }
+
+  bool IsClosed() const {
+    tensorflow::mutex_lock lock(mutex_);
+    return IsClosedInternal();
+  }
+
+  void Close() {
+    tensorflow::mutex_lock lock(mutex_);
+    reader_ = nullptr;
+    file_ = nullptr;
+  }
+
+ private:
+  bool IsClosedInternal() const {
+    return file_ == nullptr && reader_ == nullptr;
+  }
+
+  static constexpr uint64_t kReaderBufferSize = 16 * 1024 * 1024;
+
+  PyRecordRandomReader(std::unique_ptr<tensorflow::RandomAccessFile> file,
+                       std::unique_ptr<tensorflow::io::RecordReader> reader)
+      : file_(std::move(file)), reader_(std::move(reader)) {}
+
+  std::unique_ptr<tensorflow::RandomAccessFile> file_;
+  std::unique_ptr<tensorflow::io::RecordReader> reader_;
+  mutable tensorflow::mutex mutex_;
+
+  PyRecordRandomReader(const PyRecordRandomReader&) = delete;
+  void operator=(const PyRecordRandomReader&) = delete;
+};
+
+class PyRecordWriter {
+ public:
+  static absl::Status New(const std::string& filename,
+                          const tensorflow::io::RecordWriterOptions& options,
+                          PyRecordWriter** out) {
+    std::unique_ptr<tensorflow::WritableFile> file;
+    TF_RETURN_IF_ERROR(
+        tensorflow::Env::Default()->NewWritableFile(filename, &file));
+    auto writer =
+        std::make_unique<tensorflow::io::RecordWriter>(file.get(), options);
+    *out = new PyRecordWriter(std::move(file), std::move(writer));
+    return absl::OkStatus();
+  }
+
+  PyRecordWriter() = delete;
+  ~PyRecordWriter() { (void)Close(); }
+
+  absl::Status WriteRecord(absl::string_view record) {
+    tensorflow::mutex_lock lock(mutex_);
+    if (IsClosedInternal()) {
+      return absl::FailedPreconditionError("Writer is closed.");
+    }
+    return writer_->WriteRecord(record);
+  }
+
+  absl::Status Flush() {
+    tensorflow::mutex_lock lock(mutex_);
+    if (IsClosedInternal()) {
+      return absl::FailedPreconditionError("Writer is closed.");
+    }
+
+    auto status = writer_->Flush();
+    if (status.ok()) {
+      // Per the RecordWriter contract, flushing the RecordWriter does not
+      // flush the underlying file.  Here we need to do both.
+      return file_->Flush();
+    }
+    return status;
+  }
+
+  bool IsClosed() const {
+    tensorflow::mutex_lock lock(mutex_);
+    return IsClosedInternal();
+  }
+
+  absl::Status Close() {
+    tensorflow::mutex_lock lock(mutex_);
+    if (writer_ != nullptr) {
+      auto status = writer_->Close();
+      writer_ = nullptr;
+      if (!status.ok()) return status;
+    }
+    if (file_ != nullptr) {
+      auto status = file_->Close();
+      file_ = nullptr;
+      if (!status.ok()) return status;
+    }
+    return absl::OkStatus();
+  }
+
+ private:
+  bool IsClosedInternal() const { return writer_ == nullptr; }
+
+  PyRecordWriter(std::unique_ptr<tensorflow::WritableFile> file,
+                 std::unique_ptr<tensorflow::io::RecordWriter> writer)
+      : file_(std::move(file)), writer_(std::move(writer)) {}
+
+  std::unique_ptr<tensorflow::WritableFile> file_;
+  std::unique_ptr<tensorflow::io::RecordWriter> writer_;
+  mutable tensorflow::mutex mutex_;
+
+  PyRecordWriter(const PyRecordWriter&) = delete;
+  void operator=(const PyRecordWriter&) = delete;
+};
+
+PYBIND11_MODULE(_pywrap_record_io, m, pybind11::mod_gil_not_used()) {
+  py::class_<PyRecordReader>(m, "RecordIterator")
+      .def(py::init(
+          [](const std::string& filename, const std::string& compression_type) {
+            absl::Status status;
+            PyRecordReader* self = nullptr;
+            {
+              py::gil_scoped_release release;
+              status = PyRecordReader::New(filename, compression_type, &self);
+            }
+            tsl::MaybeRaiseRegisteredFromStatus(status);
+            return self;
+          }))
+      .def("__iter__", [](const py::object& self) { return self; })
+      .def("__next__",
+           [](PyRecordReader* self) {
+             bool is_closed;
+             {
+               py::gil_scoped_release release;
+               is_closed = self->IsClosed();
+             }
+             if (is_closed) {
+               throw py::stop_iteration();
+             }
+
+             tensorflow::tstring record;
+             absl::Status status;
+             {
+               py::gil_scoped_release release;
+               status = self->ReadNextRecord(&record);
+             }
+             if (absl::IsOutOfRange(status)) {
+               // Don't close because the file being read could be updated
+               // in-between
+               // __next__ calls.
+               throw py::stop_iteration();
+             }
+             tsl::MaybeRaiseRegisteredFromStatus(status);
+             return py::bytes(record);
+           })
+      .def("close",
+           [](PyRecordReader* self) {
+             py::gil_scoped_release release;
+             self->Close();
+           })
+      .def("reopen", [](PyRecordReader* self) {
+        absl::Status status;
+        {
+          py::gil_scoped_release release;
+          status = self->Reopen();
+        }
+        tsl::MaybeRaiseRegisteredFromStatus(status);
+      });
+
+  py::class_<PyRecordRandomReader>(m, "RandomRecordReader")
+      .def(py::init([](const std::string& filename) {
+        absl::Status status;
+        PyRecordRandomReader* self = nullptr;
+        {
+          py::gil_scoped_release release;
+          status = PyRecordRandomReader::New(filename, &self);
+        }
+        tsl::MaybeRaiseRegisteredFromStatus(status);
+        return self;
+      }))
+      .def("read",
+           [](PyRecordRandomReader* self, uint64_t offset) {
+             uint64_t temp_offset = offset;
+             tensorflow::tstring record;
+             absl::Status status;
+             {
+               py::gil_scoped_release release;
+               status = self->ReadRecord(&temp_offset, &record);
+             }
+             if (absl::IsOutOfRange(status)) {
+               throw py::index_error(
+                   absl::StrCat("Out of range at reading offset ", offset));
+             }
+             tsl::MaybeRaiseRegisteredFromStatus(status);
+             return py::make_tuple(py::bytes(record), temp_offset);
+           })
+      .def("close", [](PyRecordRandomReader* self) {
+        py::gil_scoped_release release;
+        self->Close();
+      });
+
+  using tensorflow::io::ZlibCompressionOptions;
+  py::class_<ZlibCompressionOptions>(m, "ZlibCompressionOptions")
+      .def_readwrite("flush_mode", &ZlibCompressionOptions::flush_mode)
+      .def_readwrite("input_buffer_size",
+                     &ZlibCompressionOptions::input_buffer_size)
+      .def_readwrite("output_buffer_size",
+                     &ZlibCompressionOptions::output_buffer_size)
+      .def_readwrite("window_bits", &ZlibCompressionOptions::window_bits)
+      .def_readwrite("compression_level",
+                     &ZlibCompressionOptions::compression_level)
+      .def_readwrite("compression_method",
+                     &ZlibCompressionOptions::compression_method)
+      .def_readwrite("mem_level", &ZlibCompressionOptions::mem_level)
+      .def_readwrite("compression_strategy",
+                     &ZlibCompressionOptions::compression_strategy);
+
+  using tensorflow::io::RecordWriterOptions;
+  py::class_<RecordWriterOptions>(m, "RecordWriterOptions")
+      .def(py::init(&RecordWriterOptions::CreateRecordWriterOptions))
+      .def_readonly("compression_type", &RecordWriterOptions::compression_type)
+      .def_readonly("zlib_options", &RecordWriterOptions::zlib_options);
+
+  using tensorflow::MaybeRaiseRegisteredFromStatus;
+
+  py::class_<PyRecordWriter>(m, "RecordWriter")
+      .def(py::init(
+          [](const std::string& filename, const RecordWriterOptions& options) {
+            PyRecordWriter* self = nullptr;
+            absl::Status status;
+            {
+              py::gil_scoped_release release;
+              status = PyRecordWriter::New(filename, options, &self);
+            }
+            tsl::MaybeRaiseRegisteredFromStatus(status);
+            return self;
+          }))
+      .def("__enter__", [](const py::object& self) { return self; })
+      .def("__exit__",
+           [](PyRecordWriter* self, py::args) {
+             absl::Status status;
+             {
+               py::gil_scoped_release release;
+               status = self->Close();
+             }
+             tsl::MaybeRaiseRegisteredFromStatus(status);
+           })
+      .def(
+          "write",
+          [](PyRecordWriter* self, absl::string_view record) {
+            absl::Status status;
+            {
+              py::gil_scoped_release release;
+              status = self->WriteRecord(record);
+            }
+            tsl::MaybeRaiseRegisteredFromStatus(status);
+          },
+          py::arg("record"))
+      .def("flush",
+           [](PyRecordWriter* self) {
+             absl::Status status;
+             {
+               py::gil_scoped_release release;
+               status = self->Flush();
+             }
+             tsl::MaybeRaiseRegisteredFromStatus(status);
+           })
+      .def("close", [](PyRecordWriter* self) {
+        absl::Status status;
+        {
+          py::gil_scoped_release release;
+          status = self->Close();
+        }
+        tsl::MaybeRaiseRegisteredFromStatus(status);
+      });
+}
+
+}  // namespace

tensorflow/python/lib/io/tf_record.py

@@ -0,0 +1,318 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+
+"""For reading and writing TFRecords files."""
+
+from tensorflow.python.lib.io import _pywrap_record_io
+from tensorflow.python.util import compat
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(
+    v1=["io.TFRecordCompressionType", "python_io.TFRecordCompressionType"])
+@deprecation.deprecated_endpoints("io.TFRecordCompressionType",
+                                  "python_io.TFRecordCompressionType")
+class TFRecordCompressionType(object):
+  """The type of compression for the record."""
+  NONE = 0
+  ZLIB = 1
+  GZIP = 2
+
+
+@tf_export(
+    "io.TFRecordOptions",
+    v1=["io.TFRecordOptions", "python_io.TFRecordOptions"])
+@deprecation.deprecated_endpoints("python_io.TFRecordOptions")
+class TFRecordOptions(object):
+  """Options used for manipulating TFRecord files."""
+  compression_type_map = {
+      TFRecordCompressionType.ZLIB: "ZLIB",
+      TFRecordCompressionType.GZIP: "GZIP",
+      TFRecordCompressionType.NONE: ""
+  }
+
+  def __init__(self,
+               compression_type=None,
+               flush_mode=None,
+               input_buffer_size=None,
+               output_buffer_size=None,
+               window_bits=None,
+               compression_level=None,
+               compression_method=None,
+               mem_level=None,
+               compression_strategy=None):
+    # pylint: disable=line-too-long
+    """Creates a `TFRecordOptions` instance.
+
+    Options only effect TFRecordWriter when compression_type is not `None`.
+    Documentation, details, and defaults can be found in
+    [`zlib_compression_options.h`](https://www.tensorflow.org/code/tensorflow/core/lib/io/zlib_compression_options.h)
+    and in the [zlib manual](http://www.zlib.net/manual.html).
+    Leaving an option as `None` allows C++ to set a reasonable default.
+
+    Args:
+      compression_type: `"GZIP"`, `"ZLIB"`, or `""` (no compression).
+      flush_mode: flush mode or `None`, Default: Z_NO_FLUSH.
+      input_buffer_size: int or `None`.
+      output_buffer_size: int or `None`.
+      window_bits: int or `None`.
+      compression_level: 0 to 9, or `None`.
+      compression_method: compression method or `None`.
+      mem_level: 1 to 9, or `None`.
+      compression_strategy: strategy or `None`. Default: Z_DEFAULT_STRATEGY.
+
+    Returns:
+      A `TFRecordOptions` object.
+
+    Raises:
+      ValueError: If compression_type is invalid.
+    """
+    # pylint: enable=line-too-long
+    # Check compression_type is valid, but for backwards compatibility don't
+    # immediately convert to a string.
+    self.get_compression_type_string(compression_type)
+    self.compression_type = compression_type
+    self.flush_mode = flush_mode
+    self.input_buffer_size = input_buffer_size
+    self.output_buffer_size = output_buffer_size
+    self.window_bits = window_bits
+    self.compression_level = compression_level
+    self.compression_method = compression_method
+    self.mem_level = mem_level
+    self.compression_strategy = compression_strategy
+
+  @classmethod
+  def get_compression_type_string(cls, options):
+    """Convert various option types to a unified string.
+
+    Args:
+      options: `TFRecordOption`, `TFRecordCompressionType`, or string.
+
+    Returns:
+      Compression type as string (e.g. `'ZLIB'`, `'GZIP'`, or `''`).
+
+    Raises:
+      ValueError: If compression_type is invalid.
+    """
+    if not options:
+      return ""
+    elif isinstance(options, TFRecordOptions):
+      return cls.get_compression_type_string(options.compression_type)
+    elif isinstance(options, TFRecordCompressionType):
+      return cls.compression_type_map[options]
+    elif options in TFRecordOptions.compression_type_map:
+      return cls.compression_type_map[options]
+    elif options in TFRecordOptions.compression_type_map.values():
+      return options
+    else:
+      raise ValueError('Not a valid compression_type: "{}"'.format(options))
+
+  def _as_record_writer_options(self):
+    """Convert to RecordWriterOptions for use with PyRecordWriter."""
+    options = _pywrap_record_io.RecordWriterOptions(
+        compat.as_bytes(
+            self.get_compression_type_string(self.compression_type)))
+
+    if self.flush_mode is not None:
+      options.zlib_options.flush_mode = self.flush_mode
+    if self.input_buffer_size is not None:
+      options.zlib_options.input_buffer_size = self.input_buffer_size
+    if self.output_buffer_size is not None:
+      options.zlib_options.output_buffer_size = self.output_buffer_size
+    if self.window_bits is not None:
+      options.zlib_options.window_bits = self.window_bits
+    if self.compression_level is not None:
+      options.zlib_options.compression_level = self.compression_level
+    if self.compression_method is not None:
+      options.zlib_options.compression_method = self.compression_method
+    if self.mem_level is not None:
+      options.zlib_options.mem_level = self.mem_level
+    if self.compression_strategy is not None:
+      options.zlib_options.compression_strategy = self.compression_strategy
+    return options
+
+
+@tf_export(v1=["io.tf_record_iterator", "python_io.tf_record_iterator"])
+@deprecation.deprecated(
+    date=None,
+    instructions=("Use eager execution and: \n"
+                  "`tf.data.TFRecordDataset(path)`"))
+def tf_record_iterator(path, options=None):
+  """An iterator that read the records from a TFRecords file.
+
+  Args:
+    path: The path to the TFRecords file.
+    options: (optional) A TFRecordOptions object.
+
+  Returns:
+    An iterator of serialized TFRecords.
+
+  Raises:
+    IOError: If `path` cannot be opened for reading.
+  """
+  compression_type = TFRecordOptions.get_compression_type_string(options)
+  return _pywrap_record_io.RecordIterator(path, compression_type)
+
+
+def tf_record_random_reader(path):
+  """Creates a reader that allows random-access reads from a TFRecords file.
+
+  The created reader object has the following method:
+
+    - `read(offset)`, which returns a tuple of `(record, ending_offset)`, where
+      `record` is the TFRecord read at the offset, and
+      `ending_offset` is the ending offset of the read record.
+
+      The method throws a `tf.errors.DataLossError` if data is corrupted at
+      the given offset. The method throws `IndexError` if the offset is out of
+      range for the TFRecords file.
+
+
+  Usage example:
+  ```py
+  reader = tf_record_random_reader(file_path)
+
+  record_1, offset_1 = reader.read(0)  # 0 is the initial offset.
+  # offset_1 is the ending offset of the 1st record and the starting offset of
+  # the next.
+
+  record_2, offset_2 = reader.read(offset_1)
+  # offset_2 is the ending offset of the 2nd record and the starting offset of
+  # the next.
+  # We can jump back and read the first record again if so desired.
+  reader.read(0)
+  ```
+
+  Args:
+    path: The path to the TFRecords file.
+
+  Returns:
+    An object that supports random-access reading of the serialized TFRecords.
+
+  Raises:
+    IOError: If `path` cannot be opened for reading.
+  """
+  return _pywrap_record_io.RandomRecordReader(path)
+
+
+@tf_export(
+    "io.TFRecordWriter", v1=["io.TFRecordWriter", "python_io.TFRecordWriter"])
+@deprecation.deprecated_endpoints("python_io.TFRecordWriter")
+class TFRecordWriter(_pywrap_record_io.RecordWriter):
+  """A class to write records to a TFRecords file.
+
+  [TFRecords tutorial](https://www.tensorflow.org/tutorials/load_data/tfrecord)
+
+  TFRecords is a binary format which is optimized for high throughput data
+  retrieval, generally in conjunction with `tf.data`. `TFRecordWriter` is used
+  to write serialized examples to a file for later consumption. The key steps
+  are:
+
+   Ahead of time:
+
+   - [Convert data into a serialized format](
+   https://www.tensorflow.org/tutorials/load_data/tfrecord#tfexample)
+   - [Write the serialized data to one or more files](
+   https://www.tensorflow.org/tutorials/load_data/tfrecord#tfrecord_files_in_python)
+
+   During training or evaluation:
+
+   - [Read serialized examples into memory](
+   https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file)
+   - [Parse (deserialize) examples](
+   https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file)
+
+  A minimal example is given below:
+
+  >>> import tempfile
+  >>> example_path = os.path.join(tempfile.gettempdir(), "example.tfrecords")
+  >>> np.random.seed(0)
+
+  >>> # Write the records to a file.
+  ... with tf.io.TFRecordWriter(example_path) as file_writer:
+  ...   for _ in range(4):
+  ...     x, y = np.random.random(), np.random.random()
+  ...
+  ...     record_bytes = tf.train.Example(features=tf.train.Features(feature={
+  ...         "x": tf.train.Feature(float_list=tf.train.FloatList(value=[x])),
+  ...         "y": tf.train.Feature(float_list=tf.train.FloatList(value=[y])),
+  ...     })).SerializeToString()
+  ...     file_writer.write(record_bytes)
+
+  >>> # Read the data back out.
+  >>> def decode_fn(record_bytes):
+  ...   return tf.io.parse_single_example(
+  ...       # Data
+  ...       record_bytes,
+  ...
+  ...       # Schema
+  ...       {"x": tf.io.FixedLenFeature([], dtype=tf.float32),
+  ...        "y": tf.io.FixedLenFeature([], dtype=tf.float32)}
+  ...   )
+
+  >>> for batch in tf.data.TFRecordDataset([example_path]).map(decode_fn):
+  ...   print("x = {x:.4f},  y = {y:.4f}".format(**batch))
+  x = 0.5488,  y = 0.7152
+  x = 0.6028,  y = 0.5449
+  x = 0.4237,  y = 0.6459
+  x = 0.4376,  y = 0.8918
+
+  This class implements `__enter__` and `__exit__`, and can be used
+  in `with` blocks like a normal file. (See the usage example above.)
+  """
+
+  # TODO(josh11b): Support appending?
+  def __init__(self, path, options=None):
+    """Opens file `path` and creates a `TFRecordWriter` writing to it.
+
+    Args:
+      path: The path to the TFRecords file.
+      options: (optional) String specifying compression type,
+          `TFRecordCompressionType`, or `TFRecordOptions` object.
+
+    Raises:
+      IOError: If `path` cannot be opened for writing.
+      ValueError: If valid compression_type can't be determined from `options`.
+    """
+    if not isinstance(options, TFRecordOptions):
+      options = TFRecordOptions(compression_type=options)
+
+    # pylint: disable=protected-access
+    super(TFRecordWriter, self).__init__(
+        compat.as_bytes(path), options._as_record_writer_options())
+    # pylint: enable=protected-access
+
+  # TODO(slebedev): The following wrapper methods are there to compensate
+  # for lack of signatures in pybind11-generated classes. Switch to
+  # __text_signature__ when TensorFlow drops Python 2.X support.
+  # See https://github.com/pybind/pybind11/issues/945
+  # pylint: disable=useless-super-delegation
+  def write(self, record):
+    """Write a string record to the file.
+
+    Args:
+      record: str
+    """
+    super(TFRecordWriter, self).write(record)
+
+  def flush(self):
+    """Flush the file."""
+    super(TFRecordWriter, self).flush()
+
+  def close(self):
+    """Close the file."""
+    super(TFRecordWriter, self).close()
+  # pylint: enable=useless-super-delegation

tensorflow/python/lib/io/tf_record_test.py

@@ -0,0 +1,646 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Tests for tf_record.TFRecordWriter and tf_record.tf_record_iterator."""
+
+import gzip
+import os
+import random
+import string
+import zlib
+
+import six
+
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.lib.io import tf_record
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+TFRecordCompressionType = tf_record.TFRecordCompressionType
+
+# Edgar Allan Poe's 'Eldorado'
+_TEXT = b"""Gaily bedight,
+    A gallant knight,
+    In sunshine and in shadow,
+    Had journeyed long,
+    Singing a song,
+    In search of Eldorado.
+
+    But he grew old
+    This knight so bold
+    And o'er his heart a shadow
+    Fell as he found
+    No spot of ground
+    That looked like Eldorado.
+
+   And, as his strength
+   Failed him at length,
+   He met a pilgrim shadow
+   'Shadow,' said he,
+   'Where can it be
+   This land of Eldorado?'
+
+   'Over the Mountains
+    Of the Moon'
+    Down the Valley of the Shadow,
+    Ride, boldly ride,'
+    The shade replied,
+    'If you seek for Eldorado!'
+    """
+
+
+class TFCompressionTestCase(test.TestCase):
+  """TFCompression Test"""
+
+  def setUp(self):
+    super(TFCompressionTestCase, self).setUp()
+    self._num_files = 2
+    self._num_records = 7
+
+  def _Record(self, f, r):
+    return compat.as_bytes("Record %d of file %d" % (r, f))
+
+  def _CreateFiles(self, options=None, prefix=""):
+    filenames = []
+    for i in range(self._num_files):
+      name = prefix + "tfrecord.%d.txt" % i
+      records = [self._Record(i, j) for j in range(self._num_records)]
+      fn = self._WriteRecordsToFile(records, name, options)
+      filenames.append(fn)
+    return filenames
+
+  def _WriteRecordsToFile(self, records, name="tfrecord", options=None):
+    fn = os.path.join(self.get_temp_dir(), name)
+    with tf_record.TFRecordWriter(fn, options=options) as writer:
+      for r in records:
+        writer.write(r)
+    return fn
+
+  def _ZlibCompressFile(self, infile, name="tfrecord.z"):
+    # zlib compress the file and write compressed contents to file.
+    with open(infile, "rb") as f:
+      cdata = zlib.compress(f.read())
+
+    zfn = os.path.join(self.get_temp_dir(), name)
+    with open(zfn, "wb") as f:
+      f.write(cdata)
+    return zfn
+
+  def _GzipCompressFile(self, infile, name="tfrecord.gz"):
+    # gzip compress the file and write compressed contents to file.
+    with open(infile, "rb") as f:
+      cdata = f.read()
+
+    gzfn = os.path.join(self.get_temp_dir(), name)
+    with gzip.GzipFile(gzfn, "wb") as f:
+      f.write(cdata)
+    return gzfn
+
+  def _ZlibDecompressFile(self, infile, name="tfrecord"):
+    with open(infile, "rb") as f:
+      cdata = zlib.decompress(f.read())
+    fn = os.path.join(self.get_temp_dir(), name)
+    with open(fn, "wb") as f:
+      f.write(cdata)
+    return fn
+
+  def _GzipDecompressFile(self, infile, name="tfrecord"):
+    with gzip.GzipFile(infile, "rb") as f:
+      cdata = f.read()
+    fn = os.path.join(self.get_temp_dir(), name)
+    with open(fn, "wb") as f:
+      f.write(cdata)
+    return fn
+
+
+class TFRecordWriterTest(TFCompressionTestCase):
+  """TFRecordWriter Test"""
+
+  def _AssertFilesEqual(self, a, b, equal):
+    for an, bn in zip(a, b):
+      with open(an, "rb") as af, open(bn, "rb") as bf:
+        if equal:
+          self.assertEqual(af.read(), bf.read())
+        else:
+          self.assertNotEqual(af.read(), bf.read())
+
+  def _CompressionSizeDelta(self, records, options_a, options_b):
+    """Validate compression with options_a and options_b and return size delta.
+
+    Compress records with options_a and options_b. Uncompress both compressed
+    files and assert that the contents match the original records. Finally
+    calculate how much smaller the file compressed with options_a was than the
+    file compressed with options_b.
+
+    Args:
+      records: The records to compress
+      options_a: First set of options to compress with, the baseline for size.
+      options_b: Second set of options to compress with.
+
+    Returns:
+      The difference in file size when using options_a vs options_b. A positive
+      value means options_a was a better compression than options_b. A negative
+      value means options_b had better compression than options_a.
+
+    """
+
+    fn_a = self._WriteRecordsToFile(records, "tfrecord_a", options=options_a)
+    test_a = list(tf_record.tf_record_iterator(fn_a, options=options_a))
+    self.assertEqual(records, test_a, options_a)
+
+    fn_b = self._WriteRecordsToFile(records, "tfrecord_b", options=options_b)
+    test_b = list(tf_record.tf_record_iterator(fn_b, options=options_b))
+    self.assertEqual(records, test_b, options_b)
+
+    # Negative number => better compression.
+    return os.path.getsize(fn_a) - os.path.getsize(fn_b)
+
+  def testWriteReadZLibFiles(self):
+    """test Write Read ZLib Files"""
+    # Write uncompressed then compress manually.
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.NONE)
+    files = self._CreateFiles(options, prefix="uncompressed")
+    zlib_files = [
+        self._ZlibCompressFile(fn, "tfrecord_%s.z" % i)
+        for i, fn in enumerate(files)
+    ]
+    self._AssertFilesEqual(files, zlib_files, False)
+
+    # Now write compressd and verify same.
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
+    compressed_files = self._CreateFiles(options, prefix="compressed")
+    self._AssertFilesEqual(compressed_files, zlib_files, True)
+
+    # Decompress compress and verify same.
+    uncompressed_files = [
+        self._ZlibDecompressFile(fn, "tfrecord_%s.z" % i)
+        for i, fn in enumerate(compressed_files)
+    ]
+    self._AssertFilesEqual(uncompressed_files, files, True)
+
+  def testWriteReadGzipFiles(self):
+    """test Write Read Gzip Files"""
+    # Write uncompressed then compress manually.
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.NONE)
+    files = self._CreateFiles(options, prefix="uncompressed")
+    gzip_files = [
+        self._GzipCompressFile(fn, "tfrecord_%s.gz" % i)
+        for i, fn in enumerate(files)
+    ]
+    self._AssertFilesEqual(files, gzip_files, False)
+
+    # Now write compressd and verify same.
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.GZIP)
+    compressed_files = self._CreateFiles(options, prefix="compressed")
+
+    # Note: Gzips written by TFRecordWriter add 'tfrecord_0' so
+    # compressed_files can't be compared with gzip_files
+
+    # Decompress compress and verify same.
+    uncompressed_files = [
+        self._GzipDecompressFile(fn, "tfrecord_%s.gz" % i)
+        for i, fn in enumerate(compressed_files)
+    ]
+    self._AssertFilesEqual(uncompressed_files, files, True)
+
+  def testNoCompressionType(self):
+    """test No Compression Type"""
+    self.assertEqual(
+        "",
+        tf_record.TFRecordOptions.get_compression_type_string(
+            tf_record.TFRecordOptions()))
+
+    self.assertEqual(
+        "",
+        tf_record.TFRecordOptions.get_compression_type_string(
+            tf_record.TFRecordOptions("")))
+
+    with self.assertRaises(ValueError):
+      tf_record.TFRecordOptions(5)
+
+    with self.assertRaises(ValueError):
+      tf_record.TFRecordOptions("BZ2")
+
+  def testZlibCompressionType(self):
+    """test Zlib Compression Type"""
+    zlib_t = tf_record.TFRecordCompressionType.ZLIB
+
+    self.assertEqual(
+        "ZLIB",
+        tf_record.TFRecordOptions.get_compression_type_string(
+            tf_record.TFRecordOptions("ZLIB")))
+
+    self.assertEqual(
+        "ZLIB",
+        tf_record.TFRecordOptions.get_compression_type_string(
+            tf_record.TFRecordOptions(zlib_t)))
+
+    self.assertEqual(
+        "ZLIB",
+        tf_record.TFRecordOptions.get_compression_type_string(
+            tf_record.TFRecordOptions(tf_record.TFRecordOptions(zlib_t))))
+
+  def testCompressionOptions(self):
+    """Create record with mix of random and repeated data to test compression on."""
+    rnd = random.Random(123)
+    random_record = compat.as_bytes(
+        "".join(rnd.choice(string.digits) for _ in range(10000)))
+    repeated_record = compat.as_bytes(_TEXT)
+    for _ in range(10000):
+      start_i = rnd.randint(0, len(_TEXT))
+      length = rnd.randint(10, 200)
+      repeated_record += _TEXT[start_i:start_i + length]
+    records = [random_record, repeated_record, random_record]
+
+    tests = [
+        ("compression_level", 2, "LE"),  # Lower compression is worse or equal.
+        ("compression_level", 6, 0),  # Default compression_level is equal.
+        ("flush_mode", zlib.Z_FULL_FLUSH, 1),  # A few less bytes.
+        ("flush_mode", zlib.Z_NO_FLUSH, 0),  # NO_FLUSH is the default.
+        ("input_buffer_size", 4096, 0),  # Increases time not size.
+        ("output_buffer_size", 4096, 0),  # Increases time not size.
+        ("window_bits", 8, -1),  # Smaller than default window increases size.
+        ("compression_strategy", zlib.Z_HUFFMAN_ONLY, -1),  # Worse.
+        ("compression_strategy", zlib.Z_FILTERED, "LE"),  # Worse or equal.
+    ]
+
+    compression_type = tf_record.TFRecordCompressionType.ZLIB
+    options_a = tf_record.TFRecordOptions(compression_type)
+    for prop, value, delta_sign in tests:
+      options_b = tf_record.TFRecordOptions(
+          compression_type=compression_type, **{prop: value})
+      delta = self._CompressionSizeDelta(records, options_a, options_b)
+      if delta_sign == "LE":
+        self.assertLessEqual(
+            delta, 0,
+            "Setting {} = {}, file was {} smaller didn't match sign of {}"
+            .format(prop, value, delta, delta_sign))
+      else:
+        self.assertTrue(
+            delta == 0 if delta_sign == 0 else delta // delta_sign > 0,
+            "Setting {} = {}, file was {} smaller didn't match sign of {}"
+            .format(prop, value, delta, delta_sign))
+
+
+class TFRecordWriterZlibTest(TFCompressionTestCase):
+  """TFRecordWriter Zlib test"""
+
+  def testZLibFlushRecord(self):
+    """test ZLib Flush Record"""
+    original = [b"small record"]
+    fn = self._WriteRecordsToFile(original, "small_record")
+    with open(fn, "rb") as h:
+      buff = h.read()
+
+    # creating more blocks and trailing blocks shouldn't break reads
+    compressor = zlib.compressobj(9, zlib.DEFLATED, zlib.MAX_WBITS)
+
+    output = b""
+    for c in buff:
+      if isinstance(c, int):
+        c = six.int2byte(c)
+      output += compressor.compress(c)
+      output += compressor.flush(zlib.Z_FULL_FLUSH)
+
+    output += compressor.flush(zlib.Z_FULL_FLUSH)
+    output += compressor.flush(zlib.Z_FULL_FLUSH)
+    output += compressor.flush(zlib.Z_FINISH)
+
+    # overwrite the original file with the compressed data
+    with open(fn, "wb") as h:
+      h.write(output)
+
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
+    actual = list(tf_record.tf_record_iterator(fn, options=options))
+    self.assertEqual(actual, original)
+
+  def testZlibReadWrite(self):
+    """Verify that files produced are zlib compatible."""
+    original = [b"foo", b"bar"]
+    fn = self._WriteRecordsToFile(original, "zlib_read_write.tfrecord")
+    zfn = self._ZlibCompressFile(fn, "zlib_read_write.tfrecord.z")
+
+    # read the compressed contents and verify.
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
+    actual = list(tf_record.tf_record_iterator(zfn, options=options))
+    self.assertEqual(actual, original)
+
+  def testZlibReadWriteLarge(self):
+    """Verify that writing large contents also works."""
+
+    # Make it large (about 5MB)
+    original = [_TEXT * 10240]
+    fn = self._WriteRecordsToFile(original, "zlib_read_write_large.tfrecord")
+    zfn = self._ZlibCompressFile(fn, "zlib_read_write_large.tfrecord.z")
+
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
+    actual = list(tf_record.tf_record_iterator(zfn, options=options))
+    self.assertEqual(actual, original)
+
+  def testGzipReadWrite(self):
+    """Verify that files produced are gzip compatible."""
+    original = [b"foo", b"bar"]
+    fn = self._WriteRecordsToFile(original, "gzip_read_write.tfrecord")
+    gzfn = self._GzipCompressFile(fn, "tfrecord.gz")
+
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.GZIP)
+    actual = list(tf_record.tf_record_iterator(gzfn, options=options))
+    self.assertEqual(actual, original)
+
+
+class TFRecordIteratorTest(TFCompressionTestCase):
+  """TFRecordIterator test"""
+
+  def setUp(self):
+    super(TFRecordIteratorTest, self).setUp()
+    self._num_records = 7
+
+  def testIterator(self):
+    """test Iterator"""
+    records = [self._Record(0, i) for i in range(self._num_records)]
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
+    fn = self._WriteRecordsToFile(records, "compressed_records", options)
+
+    reader = tf_record.tf_record_iterator(fn, options)
+    for expected in records:
+      record = next(reader)
+      self.assertEqual(expected, record)
+    with self.assertRaises(StopIteration):
+      record = next(reader)
+
+  def testWriteZlibRead(self):
+    """Verify compression with TFRecordWriter is zlib library compatible."""
+    original = [b"foo", b"bar"]
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
+    fn = self._WriteRecordsToFile(original, "write_zlib_read.tfrecord.z",
+                                  options)
+
+    zfn = self._ZlibDecompressFile(fn, "write_zlib_read.tfrecord")
+    actual = list(tf_record.tf_record_iterator(zfn))
+    self.assertEqual(actual, original)
+
+  def testWriteZlibReadLarge(self):
+    """Verify compression for large records is zlib library compatible."""
+    # Make it large (about 5MB)
+    original = [_TEXT * 10240]
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
+    fn = self._WriteRecordsToFile(original, "write_zlib_read_large.tfrecord.z",
+                                  options)
+    zfn = self._ZlibDecompressFile(fn, "write_zlib_read_large.tfrecord")
+    actual = list(tf_record.tf_record_iterator(zfn))
+    self.assertEqual(actual, original)
+
+  def testWriteGzipRead(self):
+    original = [b"foo", b"bar"]
+    options = tf_record.TFRecordOptions(TFRecordCompressionType.GZIP)
+    fn = self._WriteRecordsToFile(original, "write_gzip_read.tfrecord.gz",
+                                  options)
+
+    gzfn = self._GzipDecompressFile(fn, "write_gzip_read.tfrecord")
+    actual = list(tf_record.tf_record_iterator(gzfn))
+    self.assertEqual(actual, original)
+
+  def testReadGrowingFile_preservesReadOffset(self):
+    """Verify that tf_record_iterator preserves read offset even after EOF.
+
+    When a file is iterated to EOF, the iterator should raise StopIteration but
+    not actually close the reader. Then if later new data is appended, the
+    iterator should start returning that new data on the next call to next(),
+    preserving the read offset. This behavior is required by TensorBoard.
+    """
+    # Start the file with a good record.
+    fn = os.path.join(self.get_temp_dir(), "file.tfrecord")
+    with tf_record.TFRecordWriter(fn) as writer:
+      writer.write(b"one")
+      writer.write(b"two")
+      writer.flush()
+      iterator = tf_record.tf_record_iterator(fn)
+      self.assertEqual(b"one", next(iterator))
+      self.assertEqual(b"two", next(iterator))
+      # Iterating at EOF results in StopIteration repeatedly.
+      with self.assertRaises(StopIteration):
+        next(iterator)
+      with self.assertRaises(StopIteration):
+        next(iterator)
+      # Retrying after adding a new record successfully returns the new record,
+      # preserving the prior read offset.
+      writer.write(b"three")
+      writer.flush()
+      self.assertEqual(b"three", next(iterator))
+      with self.assertRaises(StopIteration):
+        next(iterator)
+
+  def testReadTruncatedFile_preservesReadOffset(self):
+    """Verify that tf_record_iterator throws an exception on bad TFRecords.
+
+    When a truncated record is completed, the iterator should return that new
+    record on the next attempt at iteration, preserving the read offset. This
+    behavior is required by TensorBoard.
+    """
+    # Write out a record and read it back it to get the raw bytes.
+    fn = os.path.join(self.get_temp_dir(), "temp_file")
+    with tf_record.TFRecordWriter(fn) as writer:
+      writer.write(b"truncated")
+    with open(fn, "rb") as f:
+      record_bytes = f.read()
+    # Start the file with a good record.
+    fn_truncated = os.path.join(self.get_temp_dir(), "truncated_file")
+    with tf_record.TFRecordWriter(fn_truncated) as writer:
+      writer.write(b"good")
+    with open(fn_truncated, "ab", buffering=0) as f:
+      # Cause truncation by omitting the last byte from the record.
+      f.write(record_bytes[:-1])
+      iterator = tf_record.tf_record_iterator(fn_truncated)
+      # Good record appears first.
+      self.assertEqual(b"good", next(iterator))
+      # Truncated record repeatedly causes DataLossError upon iteration.
+      with self.assertRaises(errors_impl.DataLossError):
+        next(iterator)
+      with self.assertRaises(errors_impl.DataLossError):
+        next(iterator)
+      # Retrying after completing the record successfully returns the rest of
+      # the file contents, preserving the prior read offset.
+      f.write(record_bytes[-1:])
+      self.assertEqual(b"truncated", next(iterator))
+      with self.assertRaises(StopIteration):
+        next(iterator)
+
+  def testReadReplacedFile_preservesReadOffset_afterReopen(self):
+    """Verify that tf_record_iterator allows reopening at the same read offset.
+
+    In some cases, data will be logically "appended" to a file by replacing the
+    entire file with a new version that includes the additional data. For
+    example, this can happen with certain GCS implementations (since GCS has no
+    true append operation), or when using rsync without the `--inplace` option
+    to transfer snapshots of a growing file. Since the iterator retains a handle
+    to a stale version of the file, it won't return any of the new data.
+
+    To force this to happen, callers can check for a replaced file (e.g. via a
+    stat call that reflects an increased file size) and opt to close and reopen
+    the iterator. When iteration is next attempted, this should result in
+    reading from the newly opened file, while preserving the read offset. This
+    behavior is required by TensorBoard.
+    """
+    def write_records_to_file(filename, records):
+      writer = tf_record.TFRecordWriter(filename)
+      for record in records:
+        writer.write(record)
+      writer.close()
+
+    fn = os.path.join(self.get_temp_dir(), "orig_file")
+    write_records_to_file(fn, [b"one", b"two"])
+    iterator = tf_record.tf_record_iterator(fn)
+    self.assertEqual(b"one", next(iterator))
+    self.assertEqual(b"two", next(iterator))
+    # Iterating at EOF results in StopIteration repeatedly.
+    with self.assertRaises(StopIteration):
+      next(iterator)
+    with self.assertRaises(StopIteration):
+      next(iterator)
+    # Add a new record to the end of the file by overwriting it.
+    fn2 = os.path.join(self.get_temp_dir(), "new_file")
+    write_records_to_file(fn2, [b"one", b"two", b"three"])
+    # Windows disallows replacing files while in use, so close iterator early.
+    if os.name == "nt":
+      iterator.close()
+    os.replace(fn2, fn)
+    # Iterating at EOF still results in StopIteration; new data is not shown.
+    with self.assertRaises(StopIteration):
+      next(iterator)
+    with self.assertRaises(StopIteration):
+      next(iterator)
+    # Retrying after close and reopen successfully returns the new record,
+    # preserving the prior read offset.
+    iterator.close()
+    iterator.reopen()
+    self.assertEqual(b"three", next(iterator))
+    with self.assertRaises(StopIteration):
+      next(iterator)
+
+
+class TFRecordRandomReaderTest(TFCompressionTestCase):
+
+  def testRandomReaderReadingWorks(self):
+    """Test read access to random offsets in the TFRecord file."""
+    records = [self._Record(0, i) for i in range(self._num_records)]
+    fn = self._WriteRecordsToFile(records, "uncompressed_records")
+    reader = tf_record.tf_record_random_reader(fn)
+
+    offset = 0
+    offsets = [offset]
+    # Do a pass of forward reading.
+    for i in range(self._num_records):
+      record, offset = reader.read(offset)
+      self.assertEqual(record, records[i])
+      offsets.append(offset)
+    # Reading off the bound should lead to error.
+    with self.assertRaisesRegex(IndexError, r"Out of range.*offset"):
+      reader.read(offset)
+    # Do a pass of backward reading.
+    for i in range(self._num_records - 1, 0, -1):
+      record, offset = reader.read(offsets[i])
+      self.assertEqual(offset, offsets[i + 1])
+      self.assertEqual(record, records[i])
+
+  def testRandomReaderThrowsErrorForInvalidOffset(self):
+    records = [self._Record(0, i) for i in range(self._num_records)]
+    fn = self._WriteRecordsToFile(records, "uncompressed_records")
+    reader = tf_record.tf_record_random_reader(fn)
+    with self.assertRaisesRegex(errors_impl.DataLossError, r"corrupted record"):
+      reader.read(1)  # 1 is guaranteed to be an invalid offset.
+
+  def testClosingRandomReaderCausesErrorsForFurtherReading(self):
+    records = [self._Record(0, i) for i in range(self._num_records)]
+    fn = self._WriteRecordsToFile(records, "uncompressed_records")
+    reader = tf_record.tf_record_random_reader(fn)
+    reader.close()
+    with self.assertRaisesRegex(errors_impl.FailedPreconditionError, r"closed"):
+      reader.read(0)
+
+
+class TFRecordWriterCloseAndFlushTests(test.TestCase):
+  """TFRecordWriter close and flush tests"""
+
+  # pylint: disable=arguments-differ
+  def setUp(self, compression_type=TFRecordCompressionType.NONE):
+    super(TFRecordWriterCloseAndFlushTests, self).setUp()
+    self._fn = os.path.join(self.get_temp_dir(), "tf_record_writer_test.txt")
+    self._options = tf_record.TFRecordOptions(compression_type)
+    self._writer = tf_record.TFRecordWriter(self._fn, self._options)
+    self._num_records = 20
+
+  def _Record(self, r):
+    return compat.as_bytes("Record %d" % r)
+
+  def testWriteAndLeaveOpen(self):
+    records = list(map(self._Record, range(self._num_records)))
+    for record in records:
+      self._writer.write(record)
+
+    # Verify no segfault if writer isn't explicitly closed.
+
+  def testWriteAndRead(self):
+    records = list(map(self._Record, range(self._num_records)))
+    for record in records:
+      self._writer.write(record)
+    self._writer.close()
+
+    actual = list(tf_record.tf_record_iterator(self._fn, self._options))
+    self.assertListEqual(actual, records)
+
+  def testFlushAndRead(self):
+    records = list(map(self._Record, range(self._num_records)))
+    for record in records:
+      self._writer.write(record)
+    self._writer.flush()
+
+    actual = list(tf_record.tf_record_iterator(self._fn, self._options))
+    self.assertListEqual(actual, records)
+
+  def testDoubleClose(self):
+    self._writer.write(self._Record(0))
+    self._writer.close()
+    self._writer.close()
+
+  def testFlushAfterCloseIsError(self):
+    self._writer.write(self._Record(0))
+    self._writer.close()
+
+    with self.assertRaises(errors_impl.FailedPreconditionError):
+      self._writer.flush()
+
+  def testWriteAfterCloseIsError(self):
+    self._writer.write(self._Record(0))
+    self._writer.close()
+
+    with self.assertRaises(errors_impl.FailedPreconditionError):
+      self._writer.write(self._Record(1))
+
+
+class TFRecordWriterCloseAndFlushGzipTests(TFRecordWriterCloseAndFlushTests):
+  # pylint: disable=arguments-differ
+  def setUp(self):
+    super(TFRecordWriterCloseAndFlushGzipTests,
+          self).setUp(TFRecordCompressionType.GZIP)
+
+
+class TFRecordWriterCloseAndFlushZlibTests(TFRecordWriterCloseAndFlushTests):
+  # pylint: disable=arguments-differ
+  def setUp(self):
+    super(TFRecordWriterCloseAndFlushZlibTests,
+          self).setUp(TFRecordCompressionType.ZLIB)
+
+
+if __name__ == "__main__":
+  test.main()