simplify implementation

Signed-off-by: Jean Mertz <[email protected]>
ActivityWatch · Jun 17, 2024 · 8a560fa · 8a560fa
1 parent e830395
commit 8a560fa
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 83 deletions.
diff --git a/aw-query/src/datatype.rs b/aw-query/src/datatype.rs
@@ -5,7 +5,7 @@ use std::str::FromStr as _;
 use super::functions;
 use super::QueryError;
 use aw_models::Event;
-use aw_transform::classify::{KeyValueRule, LogicalOperator, LogicalRule, RegexRule, Rule};
+use aw_transform::classify::{LogicalOperator, LogicalRule, RegexRule, Rule};
 
 use serde::{Serialize, Serializer};
 use serde_json::value::Value;
@@ -301,61 +301,42 @@ impl TryFrom<&DataType> for Rule {
 
         match rtype.as_str() {
             "none" => Ok(Self::None),
-            "or" | "and" => {
-                let Some(rules) = obj.get("rules") else {
-                    return Err(QueryError::InvalidFunctionParameters(format!(
-                        "{} rule is missing the 'rules' field",
-                        rtype
-                    )));
-                };
-
-                let rules = match rules {
-                    DataType::List(rules) => rules
-                        .iter()
-                        .map(Rule::try_from)
-                        .collect::<Result<Vec<_>, _>>()?,
-                    _ => {
-                        return Err(QueryError::InvalidFunctionParameters(format!(
-                            "the rules field of the {} rule is not a list",
-                            rtype
-                        )))
-                    }
-                };
-
-                let operator = LogicalOperator::from_str(rtype)
-                    .map_err(QueryError::InvalidFunctionParameters)?;
-
-                Ok(Rule::Logical(LogicalRule::new(rules, operator)))
-            }
+            "or" | "and" => parse_logical_rule(obj, rtype),
             "regex" => parse_regex_rule(obj),
-            "keyvalue" => {
-                let Some(rules) = obj.get("rules") else {
-                    return Err(QueryError::InvalidFunctionParameters(
-                        "keyval rule is missing the 'rules' field".to_string(),
-                    ));
-                };
-
-                let rules = match rules {
-                    DataType::Dict(rules) => rules
-                        .iter()
-                        .map(|(k, v)| Rule::try_from(v).map(|v| (k.to_owned(), v)))
-                        .collect::<Result<HashMap<_, _>, _>>()?,
-                    _ => {
-                        return Err(QueryError::InvalidFunctionParameters(
-                            "the rules field of the keyval rule is not a dict".to_string(),
-                        ))
-                    }
-                };
-
-                Ok(Rule::KeyValue(KeyValueRule::new(rules)))
-            }
             _ => Err(QueryError::InvalidFunctionParameters(format!(
                 "Unknown rule type '{rtype}'"
             ))),
         }
     }
 }
 
+fn parse_logical_rule(obj: &HashMap<String, DataType>, rtype: &String) -> Result<Rule, QueryError> {
+    let Some(rules) = obj.get("rules") else {
+        return Err(QueryError::InvalidFunctionParameters(format!(
+            "{} rule is missing the 'rules' field",
+            rtype
+        )));
+    };
+
+    let rules = match rules {
+        DataType::List(rules) => rules
+            .iter()
+            .map(Rule::try_from)
+            .collect::<Result<Vec<_>, _>>()?,
+        _ => {
+            return Err(QueryError::InvalidFunctionParameters(format!(
+                "the rules field of the {} rule is not a list",
+                rtype
+            )))
+        }
+    };
+
+    let operator =
+        LogicalOperator::from_str(rtype).map_err(QueryError::InvalidFunctionParameters)?;
+
+    Ok(Rule::Logical(LogicalRule::new(rules, operator)))
+}
+
 fn parse_regex_rule(obj: &HashMap<String, DataType>) -> Result<Rule, QueryError> {
     let regex_val = match obj.get("regex") {
         Some(regex_val) => regex_val,
@@ -385,7 +366,16 @@ fn parse_regex_rule(obj: &HashMap<String, DataType>) -> Result<Rule, QueryError>
             ))
         }
     };
-    let regex_rule = match RegexRule::new(regex_str, *ignore_case) {
+    let match_field = match obj.get("field") {
+        Some(DataType::String(v)) => Some(v.to_owned()),
+        None => None,
+        _ => {
+            return Err(QueryError::InvalidFunctionParameters(
+                "the `field` field of the regex rule is not a string".to_string(),
+            ))
+        }
+    };
+    let regex_rule = match RegexRule::new(regex_str, *ignore_case, match_field) {
         Ok(regex_rule) => regex_rule,
         Err(err) => {
             return Err(QueryError::RegexCompileError(format!(

diff --git a/aw-transform/src/classify.rs b/aw-transform/src/classify.rs
@@ -1,4 +1,4 @@
-use std::{collections::HashMap, str::FromStr};
+use std::str::FromStr;
 
 /// Transforms for classifying (tagging and categorizing) events.
 ///
@@ -10,7 +10,6 @@ pub enum Rule {
     None,
     Logical(LogicalRule),
     Regex(RegexRule),
-    KeyValue(KeyValueRule),
 }
 
 impl RuleTrait for Rule {
@@ -19,7 +18,6 @@ impl RuleTrait for Rule {
             Rule::None => false,
             Rule::Logical(rule) => rule.matches(event),
             Rule::Regex(rule) => rule.matches(event),
-            Rule::KeyValue(rule) => rule.matches(event),
         }
     }
 }
@@ -30,10 +28,15 @@ trait RuleTrait {
 
 pub struct RegexRule {
     regex: Regex,
+    field: Option<String>,
 }
 
 impl RegexRule {
-    pub fn new(regex_str: &str, ignore_case: bool) -> Result<RegexRule, fancy_regex::Error> {
+    pub fn new(
+        regex_str: &str,
+        ignore_case: bool,
+        field: Option<String>,
+    ) -> Result<RegexRule, fancy_regex::Error> {
         // can't use `RegexBuilder::case_insensitive` because it's not supported by fancy_regex,
         // so we need to prefix with `(?i)` to make it case insensitive.
         let regex = if ignore_case {
@@ -43,7 +46,7 @@ impl RegexRule {
             Regex::new(regex_str)?
         };
 
-        Ok(RegexRule { regex })
+        Ok(RegexRule { regex, field })
     }
 }
 
@@ -56,40 +59,19 @@ impl RuleTrait for RegexRule {
     fn matches(&self, event: &Event) -> bool {
         event
             .data
-            .values()
-            .filter(|val| val.is_string())
-            .any(|val| self.regex.is_match(val.as_str().unwrap()).unwrap())
+            .iter()
+            .filter(|(field, val)| {
+                self.field.as_ref().map(|v| &v == field).unwrap_or(true) && val.is_string()
+            })
+            .any(|(_, val)| self.regex.is_match(val.as_str().unwrap()).unwrap())
     }
 }
 
 impl From<Regex> for Rule {
     fn from(re: Regex) -> Self {
-        Rule::Regex(RegexRule { regex: re })
-    }
-}
-
-pub struct KeyValueRule {
-    rules: HashMap<String, Rule>,
-}
-
-impl KeyValueRule {
-    pub fn new(rules: HashMap<String, Rule>) -> Self {
-        Self { rules }
-    }
-}
-
-impl RuleTrait for KeyValueRule {
-    fn matches(&self, event: &Event) -> bool {
-        self.rules.iter().all(|(key, rule)| {
-            event
-                .data
-                .get(key)
-                .filter(|_| {
-                    let mut ev = event.clone();
-                    ev.data.retain(|k, _| k == key);
-                    rule.matches(&ev)
-                })
-                .is_some()
+        Rule::Regex(RegexRule {
+            regex: re,
+            field: None,
         })
     }
 }