aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/python/lib/io/py_record_writer.cc
diff options
context:
space:
mode:
Diffstat (limited to 'tensorflow/python/lib/io/py_record_writer.cc')
-rw-r--r--tensorflow/python/lib/io/py_record_writer.cc32
1 files changed, 22 insertions, 10 deletions
diff --git a/tensorflow/python/lib/io/py_record_writer.cc b/tensorflow/python/lib/io/py_record_writer.cc
index ba749da47a..3c64813735 100644
--- a/tensorflow/python/lib/io/py_record_writer.cc
+++ b/tensorflow/python/lib/io/py_record_writer.cc
@@ -47,6 +47,9 @@ PyRecordWriter* PyRecordWriter::New(const string& filename,
}
PyRecordWriter::~PyRecordWriter() {
+ // Writer depends on file during close for zlib flush, so destruct first.
+ writer_.reset();
+ file_.reset();
}
bool PyRecordWriter::WriteRecord(tensorflow::StringPiece record) {
@@ -56,6 +59,11 @@ bool PyRecordWriter::WriteRecord(tensorflow::StringPiece record) {
}
void PyRecordWriter::Flush(TF_Status* out_status) {
+ if (writer_ == nullptr) {
+ TF_SetStatus(out_status, TF_FAILED_PRECONDITION,
+ "Writer not initialized or previously closed");
+ return;
+ }
Status s = writer_->Flush();
if (!s.ok()) {
Set_TF_Status_from_Status(out_status, s);
@@ -64,18 +72,22 @@ void PyRecordWriter::Flush(TF_Status* out_status) {
}
void PyRecordWriter::Close(TF_Status* out_status) {
- Status s = writer_->Close();
- if (!s.ok()) {
- Set_TF_Status_from_Status(out_status, s);
- return;
+ if (writer_ != nullptr) {
+ Status s = writer_->Close();
+ if (!s.ok()) {
+ Set_TF_Status_from_Status(out_status, s);
+ return;
+ }
+ writer_.reset(nullptr);
}
- writer_.reset(nullptr);
- s = file_->Close();
- if (!s.ok()) {
- Set_TF_Status_from_Status(out_status, s);
- return;
+ if (file_ != nullptr) {
+ Status s = file_->Close();
+ if (!s.ok()) {
+ Set_TF_Status_from_Status(out_status, s);
+ return;
+ }
+ file_.reset(nullptr);
}
- file_.reset(nullptr);
}
} // namespace io