Shard that is being written.
| 50 | |
| 51 | @dataclasses.dataclass |
| 52 | class Shard(object): |
| 53 | """Shard that is being written.""" |
| 54 | |
| 55 | writer: tf.io.TFRecordWriter # TODO(sabela): use a file adapter |
| 56 | num_examples: int = 0 |
| 57 | num_bytes: int = 0 |
| 58 | |
| 59 | def add_example(self, serialized_example: str) -> None: |
| 60 | """Writes a new example.""" |
| 61 | self.writer.write(serialized_example) |
| 62 | self.num_examples += 1 |
| 63 | self.num_bytes += len(serialized_example) |
| 64 | |
| 65 | def close_writer(self) -> None: |
| 66 | """Closes the writer.""" |
| 67 | if hasattr(self.writer, 'flush'): |
| 68 | self.writer.flush() |
| 69 | self.writer.close() |
| 70 | |
| 71 | |
| 72 | @dataclasses.dataclass |