feat: added period_union transform

Author: ErikBjare

aw-models/src/event.rs

@@ -7,6 +7,7 @@ use serde_json::Map;
 use serde_json::Value;
 
 use crate::duration::DurationSerialization;
+use crate::TimeInterval;
 
 #[derive(Serialize, Deserialize, JsonSchema, Clone, Debug)]
 pub struct Event {
@@ -34,6 +35,9 @@ impl Event {
         self.timestamp
             + chrono::Duration::nanoseconds(self.duration.num_nanoseconds().unwrap() as i64)
     }
+    pub fn interval(&self) -> TimeInterval {
+        TimeInterval::new(self.timestamp, self.calculate_endtime())
+    }
 }
 
 impl PartialEq for Event {

aw-models/src/timeinterval.rs

@@ -1,3 +1,4 @@
+use std::cmp::{max, min};
 use std::fmt;
 
 use serde::de::{self, Deserialize, Deserializer, Visitor};
@@ -19,6 +20,7 @@ pub enum TimeIntervalError {
     ParseError(),
 }
 
+/// Python versions of many of these functions can be found at https://github.com/ErikBjare/timeslot
 impl TimeInterval {
     pub fn new(start: DateTime<Utc>, end: DateTime<Utc>) -> TimeInterval {
         TimeInterval { start, end }
@@ -52,6 +54,28 @@ impl TimeInterval {
     pub fn duration(&self) -> Duration {
         self.end - self.start
     }
+
+    /// If intervals are separated by a non-zero gap, return the gap as a new TimeInterval, else None
+    pub fn gap(&self, other: &TimeInterval) -> Option<TimeInterval> {
+        if self.end < other.start {
+            Some(TimeInterval::new(self.end, other.start))
+        } else if other.end < self.start {
+            Some(TimeInterval::new(other.end, self.start))
+        } else {
+            None
+        }
+    }
+
+    /// Joins two intervals together if they don't have a gap, else None
+    pub fn union(&self, other: &TimeInterval) -> Option<TimeInterval> {
+        match self.gap(other) {
+            Some(_) => None,
+            None => Some(TimeInterval::new(
+                min(self.start, other.start),
+                max(self.end, other.end),
+            )),
+        }
+    }
 }
 
 impl fmt::Display for TimeInterval {

aw-query/src/functions.rs

@@ -101,6 +101,10 @@ pub fn fill_env(env: &mut VarEnv) {
         "tag".to_string(),
         DataType::Function("tag".into(), qfunctions::tag),
     );
+    env.insert(
+        "period_union".to_string(),
+        DataType::Function("period_union".into(), qfunctions::period_union),
+    );
 }
 
 mod qfunctions {
@@ -503,6 +507,24 @@ mod qfunctions {
         }
         Ok(DataType::List(event_list))
     }
+
+    pub fn period_union(
+        args: Vec<DataType>,
+        _env: &VarEnv,
+        _ds: &Datastore,
+    ) -> Result<DataType, QueryError> {
+        // typecheck
+        validate::args_length(&args, 2)?;
+        let events1: Vec<Event> = (&args[0]).try_into()?;
+        let events2: Vec<Event> = (&args[1]).try_into()?;
+
+        let mut result = aw_transform::filter_period_intersect(&events1, &events2);
+        let mut result_tagged = Vec::new();
+        for event in result.drain(..) {
+            result_tagged.push(DataType::Event(event));
+        }
+        Ok(DataType::List(result_tagged))
+    }
 }
 
 mod validate {

aw-transform/src/lib.rs

@@ -44,3 +44,6 @@ pub use filter_period::filter_period_intersect;
 
 mod split_url;
 pub use split_url::split_url_event;
+
+mod period_union;
+pub use period_union::period_union;

aw-transform/src/period_union.rs

@@ -0,0 +1,95 @@
+use super::sort::sort_by_timestamp;
+use aw_models::Event;
+use chrono::{DateTime, Utc};
+use std::collections::VecDeque;
+
+/// Takes a list of two events and returns a new list of events covering the union
+/// of the timeperiods contained in the eventlists with no overlapping events.
+///
+/// aw-core implementation: https://github.com/ActivityWatch/aw-core/blob/b11fbe08a0405dec01380493f7b3261163cc6878/aw_transform/filter_period_intersect.py#L92
+///
+/// WARNING: This function strips all data from events as it cannot keep it consistent.
+///
+///
+/// # Example
+/// ```ignore
+///   events1   |   -------       --------- |
+///   events2   | ------  ---  --    ----   |
+///   result    | -----------  -- --------- |
+/// ```
+pub fn period_union(events1: &[Event], events2: &[Event]) -> Vec<Event> {
+    let mut sorted_events: VecDeque<Event> = VecDeque::new();
+    sorted_events.extend(sort_by_timestamp([events1, events2].concat()));
+
+    let mut events_union = Vec::new();
+
+    if !events1.is_empty() {
+        events_union.push(sorted_events.pop_front().unwrap())
+    }
+
+    for e in sorted_events {
+        let last_event = events_union.last().unwrap();
+
+        let e_p = e.interval();
+        let le_p = last_event.interval();
+
+        match e_p.union(&le_p) {
+            Some(new_period) => {
+                // If no gap and could be unioned, modify last event
+                let mut e_mod = events_union.pop().unwrap();
+                e_mod.duration = new_period.duration();
+                events_union.push(e_mod);
+            }
+            None => {
+                // If gap and could not be unioned, push event
+                events_union.push(e);
+            }
+        }
+    }
+
+    // for event in merged_events:
+    //     # Clear data
+    //     event.data = {}
+
+    events_union
+}
+
+#[cfg(test)]
+mod tests {
+    use std::str::FromStr;
+
+    use chrono::DateTime;
+    use chrono::Duration;
+    use chrono::Utc;
+    use serde_json::json;
+
+    use aw_models::Event;
+
+    use super::period_union;
+
+    #[test]
+    fn test_period_union_empty() {
+        let e_result = period_union(&[], &[]);
+        assert_eq!(e_result.len(), 0);
+    }
+
+    #[test]
+    fn test_period_union() {
+        let e1 = Event {
+            id: None,
+            timestamp: DateTime::from_str("2000-01-01T00:00:01Z").unwrap(),
+            duration: Duration::seconds(1),
+            data: json_map! {"test": json!(1)},
+        };
+
+        let mut e2 = e1.clone();
+        e2.timestamp = DateTime::from_str("2000-01-01T00:00:02Z").unwrap();
+
+        let e_result = period_union(&[e1], &[e2]);
+        assert_eq!(e_result.len(), 1);
+
+        let dt: DateTime<Utc> = DateTime::from_str("2000-01-01T00:00:01.000Z").unwrap();
+        assert_eq!(e_result[0].timestamp, dt);
+        assert_eq!(e_result[0].duration, Duration::milliseconds(2000));
+    }
+}