aw-query: Add function filter_keyvals_regex

ActivityWatch · May 31, 2020 · 132370b · 132370b
1 parent 5185b5d
commit 132370b
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 0 deletions.
diff --git a/aw-query/src/functions.rs b/aw-query/src/functions.rs
@@ -74,6 +74,13 @@ pub fn fill_env(env: &mut HashMap<String, DataType>) {
         "filter_keyvals".to_string(),
         DataType::Function("filter_keyvals".to_string(), qfunctions::filter_keyvals),
     );
+    env.insert(
+        "filter_keyvals_regex".to_string(),
+        DataType::Function(
+            "filter_keyvals_regex".to_string(),
+            qfunctions::filter_keyvals_regex,
+        ),
+    );
     env.insert(
         "filter_period_intersect".to_string(),
         DataType::Function(
@@ -420,6 +427,36 @@ mod qfunctions {
         Ok(DataType::List(filtered_tagged_events))
     }
 
+    use regex::RegexBuilder;
+
+    pub fn filter_keyvals_regex(
+        args: Vec<DataType>,
+        _env: &HashMap<String, DataType>,
+        _ds: &Datastore,
+    ) -> Result<DataType, QueryError> {
+        // typecheck
+        validate::args_length(&args, 3)?;
+        let events = (&args[0]).try_into()?;
+        let key: String = (&args[1]).try_into()?;
+        let regex_str: String = (&args[2]).try_into()?;
+        let regex = match RegexBuilder::new(&regex_str).build() {
+            Ok(regex) => regex,
+            Err(e) => {
+                return Err(QueryError::RegexCompileError(format!(
+                    "Failed to compile regex string '{}': {}",
+                    regex_str, e
+                )))
+            }
+        };
+
+        let mut filtered_events = aw_transform::filter_keyvals_regex(events, &key, &regex);
+        let mut filtered_tagged_events = Vec::new();
+        for event in filtered_events.drain(..) {
+            filtered_tagged_events.push(DataType::Event(event));
+        }
+        Ok(DataType::List(filtered_tagged_events))
+    }
+
     pub fn filter_period_intersect(
         args: Vec<DataType>,
         _env: &HashMap<String, DataType>,

diff --git a/aw-query/tests/query.rs b/aw-query/tests/query.rs
@@ -365,6 +365,7 @@ mod query_tests {
             url_events = split_url_events (events);
             filtered_events = filter_period_intersect(events, events);
             filtered_events = filter_keyvals(events, "$category", [["Uncategorized"]]);
+            filtered_events = filter_keyvals_regex(events, "key", "regex");
             chunked_events = chunk_events_by_key(events, "key");
             merged_events = merge_events_by_keys(events, ["key"]);
             return  merged_events;"#,