Merge pull request #36 from eklitzke/ctx

implement the ctx manager protocol for SummaryWriter

tensorboardX/record_writer.py

@@ -30,10 +30,6 @@ def write(self, event_str):
         w(struct.pack('I', masked_crc32c(event_str)))
         self._writer.flush()
 
-    def __del__(self):
-        if self._writer is not None:
-            self._writer.close()
-
 
 def masked_crc32c(data):
     x = u32(crc32c(data))

tensorboardX/writer.py

@@ -431,14 +431,17 @@ def add_embedding(self, mat, metadata=None, label_img=None, global_step=None, ta
         append_pbtxt(metadata, label_img, self.file_writer.get_logdir(), str(global_step).zfill(5), tag)
 
     def close(self):
+        if self.file_writer is None:
+            return  # ignore double close
         self.file_writer.flush()
         self.file_writer.close()
         for path, writer in self.all_writers.items():
             writer.flush()
             writer.close()
+        self.file_writer = self.all_writers = None
 
-    def __del__(self):
-        if self.file_writer is not None:
-            self.file_writer.close()
-        for writer in self.all_writers.values():
-            writer.close()
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()

tests/test_summary_writer.py

@@ -0,0 +1,8 @@
+from tensorboardX import SummaryWriter
+
+
+def test_summary_writer_ctx():
+    # after using a SummaryWriter as a ctx it should be closed
+    with SummaryWriter() as writer:
+        writer.add_scalar('test', 1)
+    assert writer.file_writer is None