add-keep_until

2 files changed, 39 insertions, 1 deletions

diff --git a/src/argaze.test/DataStructures.py b/src/argaze.test/DataStructures.py
index bdb166b..ca46d5e 100644
--- a/src/argaze.test/DataStructures.py
+++ b/src/argaze.test/DataStructures.py
@@ -202,6 +202,27 @@ class TestTimeStampedBufferClass(unittest.TestCase):
         self.assertEqual(len(ts_buffer), 2)
         self.assertEqual(ts_buffer.first, (2, "C"))
 
+    def test_keep_until(self):
+        """Test TimeStampedBuffer keep_until method."""
+
+        ts_buffer = DataStructures.TimeStampedBuffer({0: "A", 1: "B", 2: "C", 3: "D"})
+
+        # Check keep_until an existing timestamp
+        keep_until_2 = ts_buffer.keep_until(2)
+        
+        self.assertEqual(keep_until_2, (2, "C"))
+        self.assertEqual(len(ts_buffer), 3)
+        self.assertEqual(ts_buffer.last, (2, "C"))
+
+        # Check keep_until an none existing timestamp
+        ts_buffer = DataStructures.TimeStampedBuffer({0: "A", 1: "B", 2: "C", 3: "D"})
+
+        keep_until_1dot5 = ts_buffer.keep_until(1.5)
+
+        self.assertEqual(keep_until_1dot5, (2, "C"))
+        self.assertEqual(len(ts_buffer), 3)
+        self.assertEqual(ts_buffer.last, (2, "C"))
+
     def test_last(self):
         """Test TimeStampedBuffer last property."""
 
@@ -266,7 +287,8 @@ class TestTimeStampedBufferClass(unittest.TestCase):
         with self.assertRaises(KeyError):
 
             ts_buffer.get_last_before(-1)
-
+            
+            
     def test_get_last_until(self):
         """Test TimeStampedBuffer get_last_until method."""
 
@@ -343,6 +365,7 @@ class TestTimeStampedBufferClass(unittest.TestCase):
         self.assertEqual(ts_buffer_dataframe.columns.size, 1)
         self.assertEqual(ts_buffer_dataframe.columns[0], "value")
 
+
 if __name__ == '__main__':
 
     unittest.main()
\ No newline at end of file
diff --git a/src/argaze/DataStructures.py b/src/argaze/DataStructures.py
index a5118ff..437cb93 100644
--- a/src/argaze/DataStructures.py
+++ b/src/argaze/DataStructures.py
@@ -144,6 +144,20 @@ class TimeStampedBuffer(collections.OrderedDict):
 
         return popep_ts, poped_value
 
+    def keep_until(self, ts: TimeStampType) -> Tuple[TimeStampType, DataType]:
+        """Pop all item after a given timestamped value and return the first popped value."""
+
+        # get first item after given timestamp
+        last_kept_ts, last_kept_value = self.get_first_from(ts)
+
+        last_ts, last_value = self.last
+
+        while last_kept_ts < last_ts:
+            self.pop_last()
+            last_ts, last_value = self.last
+            
+        return last_ts, last_value
+
     @property
     def last(self) -> Tuple[TimeStampType, DataType]:
         """Easing access to last item."""
@@ -186,6 +200,7 @@ class TimeStampedBuffer(collections.OrderedDict):
         else:
             
             raise KeyError(f'No data stored before {ts} timestamp.')
+        
 
     def get_last_until(self, ts) -> Tuple[TimeStampType, DataType]:
         """Retreive last item timestamp until a given timestamp value."""