1use serde::{Deserialize, Serialize};
22use std::collections::{HashMap, HashSet};
23
24#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
33#[serde(untagged)]
34pub enum AttributeValue {
35 Bool(bool),
36 Int(i64),
37 Float(f64),
38 String(String),
39 StringSet(HashSet<String>),
40}
41
42impl From<bool> for AttributeValue {
43 fn from(v: bool) -> Self {
44 AttributeValue::Bool(v)
45 }
46}
47impl From<i64> for AttributeValue {
48 fn from(v: i64) -> Self {
49 AttributeValue::Int(v)
50 }
51}
52impl From<f64> for AttributeValue {
53 fn from(v: f64) -> Self {
54 AttributeValue::Float(v)
55 }
56}
57impl From<&str> for AttributeValue {
58 fn from(v: &str) -> Self {
59 AttributeValue::String(v.to_string())
60 }
61}
62impl From<String> for AttributeValue {
63 fn from(v: String) -> Self {
64 AttributeValue::String(v)
65 }
66}
67impl From<HashSet<String>> for AttributeValue {
68 fn from(v: HashSet<String>) -> Self {
69 AttributeValue::StringSet(v)
70 }
71}
72
73#[derive(Debug, Clone, Default, Serialize, Deserialize)]
80pub struct AttributeBag {
81 attrs: HashMap<String, AttributeValue>,
82}
83
84impl AttributeBag {
85 pub fn new() -> Self {
86 Self {
87 attrs: HashMap::new(),
88 }
89 }
90
91 pub fn set(&mut self, key: impl Into<String>, value: impl Into<AttributeValue>) {
92 self.attrs.insert(key.into(), value.into());
93 }
94
95 pub fn get(&self, key: &str) -> Option<&AttributeValue> {
96 self.attrs.get(key)
97 }
98
99 pub fn contains(&self, key: &str) -> bool {
100 self.attrs.contains_key(key)
101 }
102
103 pub fn get_bool(&self, key: &str) -> Option<bool> {
104 match self.get(key) {
105 Some(AttributeValue::Bool(v)) => Some(*v),
106 _ => None,
107 }
108 }
109
110 pub fn get_int(&self, key: &str) -> Option<i64> {
111 match self.get(key) {
112 Some(AttributeValue::Int(v)) => Some(*v),
113 _ => None,
114 }
115 }
116
117 pub fn get_float(&self, key: &str) -> Option<f64> {
118 match self.get(key) {
119 Some(AttributeValue::Float(v)) => Some(*v),
120 Some(AttributeValue::Int(v)) => Some(*v as f64),
123 _ => None,
124 }
125 }
126
127 pub fn get_string(&self, key: &str) -> Option<&str> {
128 match self.get(key) {
129 Some(AttributeValue::String(v)) => Some(v.as_str()),
130 _ => None,
131 }
132 }
133
134 pub fn get_string_set(&self, key: &str) -> Option<&HashSet<String>> {
135 match self.get(key) {
136 Some(AttributeValue::StringSet(v)) => Some(v),
137 _ => None,
138 }
139 }
140
141 pub fn set_contains(&self, key: &str, value: &str) -> bool {
143 self.get_string_set(key)
144 .map(|set| set.contains(value))
145 .unwrap_or(false)
146 }
147
148 pub fn len(&self) -> usize {
149 self.attrs.len()
150 }
151
152 pub fn is_empty(&self) -> bool {
153 self.attrs.is_empty()
154 }
155
156 pub fn iter(&self) -> impl Iterator<Item = (&str, &AttributeValue)> {
157 self.attrs.iter().map(|(k, v)| (k.as_str(), v))
158 }
159}
160
161pub trait AttributeExtractor {
172 fn extract(&self, bag: &mut AttributeBag);
173}
174
175#[cfg(test)]
176mod tests {
177 use super::*;
178
179 #[test]
180 fn basic_bag() {
181 let mut bag = AttributeBag::new();
182 bag.set("authenticated", true);
183 bag.set("delegation.depth", 2i64);
184 bag.set("subject.id", "alice@corp.com");
185 bag.set("intent.confidence", 0.92f64);
186
187 assert_eq!(bag.get_bool("authenticated"), Some(true));
188 assert_eq!(bag.get_int("delegation.depth"), Some(2));
189 assert_eq!(bag.get_string("subject.id"), Some("alice@corp.com"));
190 assert_eq!(bag.get_float("intent.confidence"), Some(0.92));
191 }
192
193 #[test]
194 fn int_to_float_promotion() {
195 let mut bag = AttributeBag::new();
196 bag.set("delegation.depth", 2i64);
197 assert_eq!(bag.get_float("delegation.depth"), Some(2.0));
198 }
199
200 #[test]
201 fn string_set_contains() {
202 let mut bag = AttributeBag::new();
203 bag.set(
204 "session.labels",
205 HashSet::from(["PII".to_string(), "financial".to_string()]),
206 );
207
208 assert!(bag.set_contains("session.labels", "PII"));
209 assert!(bag.set_contains("session.labels", "financial"));
210 assert!(!bag.set_contains("session.labels", "PHI"));
211 }
212
213 #[test]
214 fn missing_keys() {
215 let bag = AttributeBag::new();
216 assert_eq!(bag.get_bool("nonexistent"), None);
217 assert_eq!(bag.get_int("nonexistent"), None);
218 assert!(!bag.set_contains("nonexistent", "value"));
219 }
220
221 #[test]
222 fn type_mismatch_returns_none() {
223 let mut bag = AttributeBag::new();
224 bag.set("subject.id", "alice");
225 assert_eq!(bag.get_bool("subject.id"), None);
227 assert_eq!(bag.get_int("subject.id"), None);
228 }
229}