use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum AttributeValue {
Bool(bool),
Int(i64),
Float(f64),
String(String),
StringSet(HashSet<String>),
}
impl From<bool> for AttributeValue {
fn from(v: bool) -> Self {
AttributeValue::Bool(v)
}
}
impl From<i64> for AttributeValue {
fn from(v: i64) -> Self {
AttributeValue::Int(v)
}
}
impl From<f64> for AttributeValue {
fn from(v: f64) -> Self {
AttributeValue::Float(v)
}
}
impl From<&str> for AttributeValue {
fn from(v: &str) -> Self {
AttributeValue::String(v.to_string())
}
}
impl From<String> for AttributeValue {
fn from(v: String) -> Self {
AttributeValue::String(v)
}
}
impl From<HashSet<String>> for AttributeValue {
fn from(v: HashSet<String>) -> Self {
AttributeValue::StringSet(v)
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AttributeBag {
attrs: HashMap<String, AttributeValue>,
}
impl AttributeBag {
pub fn new() -> Self {
Self {
attrs: HashMap::new(),
}
}
pub fn set(&mut self, key: impl Into<String>, value: impl Into<AttributeValue>) {
self.attrs.insert(key.into(), value.into());
}
pub fn get(&self, key: &str) -> Option<&AttributeValue> {
self.attrs.get(key)
}
pub fn contains(&self, key: &str) -> bool {
self.attrs.contains_key(key)
}
pub fn get_bool(&self, key: &str) -> Option<bool> {
match self.get(key) {
Some(AttributeValue::Bool(v)) => Some(*v),
_ => None,
}
}
pub fn get_int(&self, key: &str) -> Option<i64> {
match self.get(key) {
Some(AttributeValue::Int(v)) => Some(*v),
_ => None,
}
}
pub fn get_float(&self, key: &str) -> Option<f64> {
match self.get(key) {
Some(AttributeValue::Float(v)) => Some(*v),
Some(AttributeValue::Int(v)) => Some(*v as f64),
_ => None,
}
}
pub fn get_string(&self, key: &str) -> Option<&str> {
match self.get(key) {
Some(AttributeValue::String(v)) => Some(v.as_str()),
_ => None,
}
}
pub fn get_string_set(&self, key: &str) -> Option<&HashSet<String>> {
match self.get(key) {
Some(AttributeValue::StringSet(v)) => Some(v),
_ => None,
}
}
pub fn set_contains(&self, key: &str, value: &str) -> bool {
self.get_string_set(key)
.map(|set| set.contains(value))
.unwrap_or(false)
}
pub fn len(&self) -> usize {
self.attrs.len()
}
pub fn is_empty(&self) -> bool {
self.attrs.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = (&str, &AttributeValue)> {
self.attrs.iter().map(|(k, v)| (k.as_str(), v))
}
}
pub trait AttributeExtractor {
fn extract(&self, bag: &mut AttributeBag);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn basic_bag() {
let mut bag = AttributeBag::new();
bag.set("authenticated", true);
bag.set("delegation.depth", 2i64);
bag.set("subject.id", "alice@corp.com");
bag.set("intent.confidence", 0.92f64);
assert_eq!(bag.get_bool("authenticated"), Some(true));
assert_eq!(bag.get_int("delegation.depth"), Some(2));
assert_eq!(bag.get_string("subject.id"), Some("alice@corp.com"));
assert_eq!(bag.get_float("intent.confidence"), Some(0.92));
}
#[test]
fn int_to_float_promotion() {
let mut bag = AttributeBag::new();
bag.set("delegation.depth", 2i64);
assert_eq!(bag.get_float("delegation.depth"), Some(2.0));
}
#[test]
fn string_set_contains() {
let mut bag = AttributeBag::new();
bag.set(
"session.labels",
HashSet::from(["PII".to_string(), "financial".to_string()]),
);
assert!(bag.set_contains("session.labels", "PII"));
assert!(bag.set_contains("session.labels", "financial"));
assert!(!bag.set_contains("session.labels", "PHI"));
}
#[test]
fn missing_keys() {
let bag = AttributeBag::new();
assert_eq!(bag.get_bool("nonexistent"), None);
assert_eq!(bag.get_int("nonexistent"), None);
assert!(!bag.set_contains("nonexistent", "value"));
}
#[test]
fn type_mismatch_returns_none() {
let mut bag = AttributeBag::new();
bag.set("subject.id", "alice");
assert_eq!(bag.get_bool("subject.id"), None);
assert_eq!(bag.get_int("subject.id"), None);
}
}