Skip to main content

wpl/ast/
ann_func.rs

1use crate::ast::AnnFun;
2use crate::{WparseError, WparseReason};
3use orion_error::{ToStructError, UvsFrom};
4use smol_str::SmolStr;
5use std::collections::BTreeMap;
6use wp_connector_api::SourceEvent;
7use wp_model_core::model::{DataField, DataRecord};
8use wp_model_core::raw::RawData;
9
10pub trait AnnotationFunc {
11    fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError>;
12}
13
14#[derive(Clone, Debug)]
15pub struct TagAnnotation {
16    args: BTreeMap<SmolStr, SmolStr>,
17}
18
19impl AnnotationFunc for TagAnnotation {
20    fn proc(&self, _src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
21        for (key, val) in &self.args {
22            data.append(DataField::from_chars(key.clone(), val.clone()));
23        }
24        Ok(())
25    }
26}
27
28#[derive(Clone, Debug)]
29pub struct NoopAnnotation;
30
31impl AnnotationFunc for NoopAnnotation {
32    fn proc(&self, _src: &SourceEvent, _data: &mut DataRecord) -> Result<(), WparseError> {
33        Ok(())
34    }
35}
36
37#[derive(Clone, Debug)]
38pub struct RawCopy {
39    raw_key: SmolStr,
40}
41
42impl AnnotationFunc for RawCopy {
43    fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
44        match &src.payload {
45            RawData::String(raw) => {
46                data.append(DataField::from_chars(self.raw_key.clone(), raw.clone()));
47            }
48            RawData::Bytes(raw) => match std::str::from_utf8(raw) {
49                Ok(str) => {
50                    data.append(DataField::from_chars(self.raw_key.clone(), str.to_string()));
51                }
52                Err(e) => {
53                    return Err(WparseReason::from_data()
54                        .to_err()
55                        .with_detail(format!("[u8] to string error :{}", e)));
56                }
57            },
58            RawData::ArcBytes(raw) => match std::str::from_utf8(raw) {
59                Ok(str) => {
60                    data.append(DataField::from_chars(self.raw_key.clone(), str.to_string()));
61                }
62                Err(e) => {
63                    return Err(WparseReason::from_data()
64                        .to_err()
65                        .with_detail(format!("ArcBytes to string error :{}", e)));
66                }
67            },
68        }
69        Ok(())
70    }
71}
72
73#[derive(Clone, Debug)]
74pub enum AnnotationType {
75    Tag(TagAnnotation),
76    Copy(RawCopy),
77    Null(NoopAnnotation),
78}
79
80impl AnnotationFunc for AnnotationType {
81    fn proc(&self, src: &SourceEvent, data: &mut DataRecord) -> Result<(), WparseError> {
82        match self {
83            AnnotationType::Tag(func) => func.proc(src, data),
84            AnnotationType::Null(func) => func.proc(src, data),
85            AnnotationType::Copy(func) => func.proc(src, data),
86        }
87    }
88}
89
90impl AnnotationType {
91    pub fn convert(ann: &Option<AnnFun>) -> Vec<Self> {
92        let mut vec = vec![];
93        if let Some(ann) = ann {
94            if !ann.tags.is_empty() {
95                vec.push(AnnotationType::Tag(TagAnnotation {
96                    args: ann.tags.clone(),
97                }));
98            }
99
100            if let Some((k, v)) = &ann.copy_raw {
101                if k == "name" {
102                    vec.push(AnnotationType::Copy(RawCopy { raw_key: v.clone() }));
103                } else {
104                    vec.push(AnnotationType::Null(NoopAnnotation {}))
105                }
106            }
107        } else {
108            vec.push(AnnotationType::Null(NoopAnnotation {}))
109        }
110        vec
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    use crate::pkg::DEFAULT_KEY;
118    use orion_error::TestAssert;
119    use std::collections::BTreeMap;
120    use wp_connector_api::{SourceEvent, Tags};
121    use wp_model_core::model::DataRecord;
122    use wp_model_core::raw::RawData;
123
124    #[test]
125    fn test_tag_fun() {
126        let ann = AnnFun {
127            tags: BTreeMap::from([("tag_1".into(), "x".into())]),
128            copy_raw: None,
129        };
130        let tag = AnnotationType::convert(&Some(ann));
131        let mut data = DataRecord::test_value();
132        let src = SourceEvent::new(
133            1,
134            DEFAULT_KEY.to_string(),
135            RawData::String("test".to_string()),
136            Tags::new().into(),
137        );
138        tag.first().unwrap().proc(&src, &mut data).assert();
139        let expected = DataField::from_chars("tag_1", "x");
140        assert_eq!(data.field("tag_1").map(|s| s.as_field()), Some(&expected));
141    }
142
143    #[test]
144    fn test_copy_fun() {
145        let ann = AnnFun {
146            tags: Default::default(),
147            copy_raw: Some(("name".into(), "raw".into())),
148        };
149        let tag = AnnotationType::convert(&Some(ann));
150        let mut data = DataRecord::test_value();
151        let src = SourceEvent::new(
152            1,
153            DEFAULT_KEY.to_string(),
154            RawData::String("test".to_string()),
155            Tags::new().into(),
156        );
157        tag.first().unwrap().proc(&src, &mut data).unwrap();
158        let expected = DataField::from_chars("raw", "test");
159        assert_eq!(data.field("raw").map(|s| s.as_field()), Some(&expected));
160    }
161}