Enforcing data type

2 files changed, 26 insertions, 3 deletions

diff --git a/src/argaze/DataAnalysis/GazeAnalysis.py b/src/argaze/DataAnalysis/GazeAnalysis.py
index 2968b2a..d0e4b9e 100644
--- a/src/argaze/DataAnalysis/GazeAnalysis.py
+++ b/src/argaze/DataAnalysis/GazeAnalysis.py
@@ -15,6 +15,16 @@ class GazePosition(DictObject.DictObject):
 
         super().__init__(type(self).__name__, **{'x': x, 'y': y})
 
+class TimeStampedGazePositionBuffer(TimeStampedDataBuffer.TimeStampedDataBuffer):
+    """Define timestamped data buffer for gaze position only"""
+
+    def __setitem__(self, key, value: GazePosition):
+        """Force value to be a GazePosition"""
+        if type(value) != GazePosition:
+            raise ValueError('value must be a GazePosition')
+
+        super().__setitem__(key, value)
+
 class Fixation(DictObject.DictObject):
     """Define fixation data"""
 
@@ -22,15 +32,28 @@ class Fixation(DictObject.DictObject):
 
         super().__init__(type(self).__name__, **{'duration': duration, 'dispersion': dispersion, 'centroid': [cx, cy]})
 
+class TimeStampedFixationBuffer(TimeStampedDataBuffer.TimeStampedDataBuffer):
+    """Define timestamped data buffer for fixation only"""
+
+    def __setitem__(self, key, value: Fixation):
+        """Force value to be a Fixation"""
+        if type(value) != Fixation:
+            raise ValueError('value must be a Fixation')
+
+        super().__setitem__(key, value)
+
 class FixationAnalyser():
 
-    fixations = TimeStampedDataBuffer.TimeStampedDataBuffer()
+    fixations = TimeStampedFixationBuffer()
     saccades = TimeStampedDataBuffer.TimeStampedDataBuffer()
 
     def analyse(self, ts_gaze_position_buffer):
         raise NotImplementedError('analyse() method not implemented')
 
-    def __init__(self, ts_gaze_position_buffer):
+    def __init__(self, ts_gaze_position_buffer: TimeStampedGazePositionBuffer):
+
+        if type(ts_gaze_position_buffer) != TimeStampedGazePositionBuffer:
+            raise ValueError('argument must be a TimeStampedGazePositionBuffer')
 
         # do analysis on a copy
         self.analyse(ts_gaze_position_buffer.copy())
diff --git a/src/argaze/utils/analyse_tobii_segment_fixations.py b/src/argaze/utils/analyse_tobii_segment_fixations.py
index dcd38ea..a07a515 100644
--- a/src/argaze/utils/analyse_tobii_segment_fixations.py
+++ b/src/argaze/utils/analyse_tobii_segment_fixations.py
@@ -32,7 +32,7 @@ def main():
         print(f'{len(tobii_ts_gaze_position_buffer)} gaze positions loaded')
 
         # format tobii gaze data into generic gaze data
-        generic_ts_gaze_position_buffer = TimeStampedDataBuffer.TimeStampedDataBuffer()
+        generic_ts_gaze_position_buffer = GazeAnalysis.TimeStampedGazePositionBuffer()
 
         for ts, tobii_data in tobii_ts_gaze_position_buffer.items():
             generic_data = GazeAnalysis.GazePosition(tobii_data.gp[0] * tobii_segment_video.get_width(), tobii_data.gp[1] * tobii_segment_video.get_height())