From fd73732c6acdbea3ecf00c315d4a40d38864bb99 Mon Sep 17 00:00:00 2001
From: Joel Collins <joel@jtcollins.net>
Date: Tue, 6 Nov 2018 18:37:17 +0000
Subject: [PATCH] StreamingCamera cleans up StreamObject files automatically on
 close()

---
 openflexure_microscope/camera/base.py |  8 ++++++--
 openflexure_microscope/camera/pi.py   | 17 +++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/openflexure_microscope/camera/base.py b/openflexure_microscope/camera/base.py
index 9b561117..5752db6c 100644
--- a/openflexure_microscope/camera/base.py
+++ b/openflexure_microscope/camera/base.py
@@ -55,6 +55,9 @@ class StreamObject(object):
         # Keep on disk after close by default
         self.keep_on_disk = True
 
+        # Log if created by context manager
+        self.context_manager = False
+
         # Object lock
         self.locked = False  # Is the StreamObject locked for writing? (Handled by StreamingCamera)
 
@@ -62,6 +65,7 @@ class StreamObject(object):
         log("Entering context for {}.\
             Stored files will be cleaned up automatically.".format(self.id))
         self.keep_on_disk = False  # Flag file to be removed on close.
+        self.context_manager = True  # Used in metadata
         return self
 
     def __exit__(self, *args):
@@ -71,7 +75,7 @@ class StreamObject(object):
     def build_file_path(self, filename: str, folder: str, fmt: str, iterator: int=0) -> str:
         """
         Construct a full file path, based on filename, folder, and file format. 
-        
+
         Defaults to datestamp. Iterator adds a numeric increment to the file name.
         """
         if filename:
@@ -128,6 +132,7 @@ class StreamObject(object):
             'locked': self.locked,
             'keep_on_disk': self.keep_on_disk,
             'path': self.file,
+            'context_manager': self.context_manager,
         }
 
         # Get file path
@@ -211,7 +216,6 @@ class StreamObject(object):
             os.remove(self.file)
             return True
         else:
-            log("File not found {}".format(self.file))
             return False
 
     def shunt(self):
diff --git a/openflexure_microscope/camera/pi.py b/openflexure_microscope/camera/pi.py
index eb868c04..27029bb9 100644
--- a/openflexure_microscope/camera/pi.py
+++ b/openflexure_microscope/camera/pi.py
@@ -96,6 +96,23 @@ class StreamingCamera(BaseCamera):
     def initialisation(self):
         """Run any initialisation code when the frame iterator starts."""
         pass
+    
+    # HANDLE CONTEXT MANAGER AND FILE CLOSING
+    
+    def close(self):
+        # Close all StreamObjects
+        for capture_list in [self.images, self.videos]:
+            for stream_object in capture_list:
+                stream_object.close()
+
+    def __enter__(self):
+        log("Entering context for {}.\
+            Stored files will be cleaned up automatically.".format(self))
+        return self
+
+    def __exit__(self, *args):
+        log("Cleaning up {}".format(self))
+        self.close()
 
     # RETURNING CAPTURES