From 8bfb8ae65205cebce16c6557d9e271dc7e2babad Mon Sep 17 00:00:00 2001
From: Samuel Gaist <samuel.gaist@idiap.ch>
Date: Thu, 14 May 2020 15:26:31 +0200
Subject: [PATCH] [data] Add reset method to DataSource classes

This allows to reset the state of a data source.

Currently only used by the CachedDataSource
---
 beat/backend/python/data.py           | 19 +++++++++++++++++--
 beat/backend/python/test/test_data.py | 12 ++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/beat/backend/python/data.py b/beat/backend/python/data.py
index ad30337..c126d55 100644
--- a/beat/backend/python/data.py
+++ b/beat/backend/python/data.py
@@ -209,6 +209,14 @@ class DataSource(object):
     def close(self):
         self.infos = []
 
+    def reset(self):
+        """Reset the state of the data source
+
+        This shall only clear the current state, not require a new call
+        to setup the source.
+        """
+        pass
+
     def __del__(self):
         """Makes sure all resources are released when the object is deleted"""
         self.close()
@@ -467,11 +475,18 @@ class CachedDataSource(DataSource):
         return True
 
     def close(self):
-        if self.current_file is not None:
-            self.current_file.close()
+        self.reset()
 
         super(CachedDataSource, self).close()
 
+    def reset(self):
+        """Rest the current state"""
+
+        if self.current_file is not None:
+            self.current_file.close()
+            self.current_file = None
+        self.current_file_index = None
+
     def __getitem__(self, index):
         """Retrieve a block of data
 
diff --git a/beat/backend/python/test/test_data.py b/beat/backend/python/test/test_data.py
index f24d89e..227814a 100644
--- a/beat/backend/python/test/test_data.py
+++ b/beat/backend/python/test/test_data.py
@@ -348,6 +348,18 @@ class TestCachedDataSource(TestCachedDataBase):
         self.check_valid_data_indices(cached_file)
         self.check_invalid_indices(cached_file)
 
+    def test_reset(self):
+        self.writeData("user/single_integer/1", 0, 9)
+
+        cached_source = CachedDataSource()
+        cached_source.setup(self.filename, prefix)
+        _, _, _ = cached_source[0]
+        self.assertIsNotNone(cached_source.current_file_index)
+        self.assertIsNotNone(cached_source.current_file)
+        cached_source.reset()
+        self.assertIsNone(cached_source.current_file_index)
+        self.assertIsNone(cached_source.current_file)
+
 
 # ----------------------------------------------------------
 
-- 
2.21.0