blob: e557756cbcce00b9ed6e70be8e5252c713177dc5 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
|
#include "tensorflow/python/lib/io/py_record_writer.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/platform/port.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/public/env.h"
namespace tensorflow {
namespace io {
PyRecordWriter::PyRecordWriter() {}
PyRecordWriter* PyRecordWriter::New(const string& filename) {
WritableFile* file;
Status s = Env::Default()->NewWritableFile(filename, &file);
if (!s.ok()) {
return nullptr;
}
PyRecordWriter* writer = new PyRecordWriter;
writer->file_ = file;
writer->writer_ = new RecordWriter(writer->file_);
return writer;
}
PyRecordWriter::~PyRecordWriter() {
delete writer_;
delete file_;
}
bool PyRecordWriter::WriteRecord(::tensorflow::StringPiece record) {
if (writer_ == nullptr) return false;
Status s = writer_->WriteRecord(record);
return s.ok();
}
void PyRecordWriter::Close() {
delete writer_;
delete file_;
writer_ = nullptr;
file_ = nullptr;
}
} // namespace io
} // namespace tensorflow
|
