[tf.data] Replace `__builtins__.open` in `tf.data.experimental.{save,load}` with `tf.io.gfile.GFile`.
PiperOrigin-RevId: 357057324
Change-Id: I7e96eb56404acf09a170c4084dd6bfe5005e2461
diff --git a/tensorflow/python/data/experimental/ops/io.py b/tensorflow/python/data/experimental/ops/io.py
index b91899a..d102f17 100644
--- a/tensorflow/python/data/experimental/ops/io.py
+++ b/tensorflow/python/data/experimental/ops/io.py
@@ -25,6 +25,7 @@
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import gen_experimental_dataset_ops
+from tensorflow.python.platform import gfile
from tensorflow.python.util import lazy_loader
from tensorflow.python.util.tf_export import tf_export
@@ -98,8 +99,8 @@
coder = nested_structure_coder.StructureCoder()
encoded = coder.encode_structure(dataset.element_spec)
- os.makedirs(path, exist_ok=True)
- with open(os.path.join(path, DATASET_SPEC_FILENAME), "wb") as f:
+ gfile.MakeDirs(path)
+ with gfile.GFile(os.path.join(path, DATASET_SPEC_FILENAME), "wb") as f:
f.write(encoded.SerializeToString())
path = ops.convert_to_tensor(path, dtype=dtypes.string, name="path")
@@ -131,7 +132,7 @@
self._path = path
if element_spec is None:
- with open(os.path.join(path, DATASET_SPEC_FILENAME), "rb") as f:
+ with gfile.GFile(os.path.join(path, DATASET_SPEC_FILENAME), "rb") as f:
encoded_spec = f.read()
struct_pb = nested_structure_coder.struct_pb2.StructuredValue()
struct_pb.ParseFromString(encoded_spec)
