Add close and context manager

km3io/offline.py

@@ -323,8 +323,18 @@ class OfflineReader:
 
         """
         self._fobj = uproot.open(file_path)
+        self._filename = file_path
         self._tree = self._fobj[MAIN_TREE_NAME]
 
+    def close(self):
+        self._fobj.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
     @cached_property
     def events(self):
         """The `E` branch, containing all offline events."""

km3io/online.py

@@ -111,10 +111,20 @@ class OnlineReader:
     """Reader for online ROOT files"""
     def __init__(self, filename):
         self._fobj = uproot.open(filename)
+        self._filename = filename
         self._events = None
         self._timeslices = None
         self._summaryslices = None
 
+    def close(self):
+        self._fobj.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
     @property
     def events(self):
         if self._events is None:

tests/test_offline.py

@@ -77,6 +77,11 @@ class TestOfflineReader(unittest.TestCase):
         self.nu = OFFLINE_NUMUCC
         self.n_events = 10
 
+    def test_context_manager(self):
+        filename = SAMPLES_DIR / 'aanet_v2.0.0.root'
+        with OfflineReader(filename) as r:
+            assert r._filename == filename
+
     def test_number_events(self):
         assert self.n_events == len(self.r.events)
 

tests/test_online.py

@@ -7,6 +7,13 @@ from km3io.online import OnlineReader, get_rate, has_udp_trailer, get_udp_max_se
 SAMPLES_DIR = os.path.join(os.path.dirname(__file__), "samples")
 
 
+class TestOnlineReaderContextManager(unittest.TestCase):
+    def test_context_manager(self):
+        filename = os.path.join(SAMPLES_DIR, "daq_v1.0.0.root")
+        with OnlineReader(filename) as r:
+            assert r._filename == filename
+
+
 class TestOnlineEvents(unittest.TestCase):
     def setUp(self):
         self.events = OnlineReader(os.path.join(SAMPLES_DIR,