1extern crate alloc;
4use alloc::collections::BTreeMap;
5use alloc::string::String;
6use alloc::vec::Vec;
7use core::fmt;
8use core::str::FromStr;
9
10use crate::{EdgeRelation, Header, Id128, Timestamp};
11
12#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
19#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
20#[cfg_attr(feature = "serde", serde(rename_all = "snake_case"))]
21pub enum EntityKind {
22 #[default]
25 Concept,
26 Document,
29 Dataset,
32 Project,
35 Person,
37 Org,
39}
40
41impl EntityKind {
42 pub const ALL: [Self; 6] = [
43 Self::Concept,
44 Self::Document,
45 Self::Dataset,
46 Self::Project,
47 Self::Person,
48 Self::Org,
49 ];
50
51 pub const fn name(self) -> &'static str {
52 match self {
53 Self::Concept => "concept",
54 Self::Document => "document",
55 Self::Dataset => "dataset",
56 Self::Project => "project",
57 Self::Person => "person",
58 Self::Org => "org",
59 }
60 }
61}
62
63impl fmt::Display for EntityKind {
64 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65 f.write_str(self.name())
66 }
67}
68
69const ENTITY_KIND_VALID: &[&str] = &["concept", "document", "dataset", "project", "person", "org"];
70
71impl FromStr for EntityKind {
72 type Err = crate::error::UnknownVariant;
73
74 fn from_str(s: &str) -> Result<Self, Self::Err> {
75 match s.trim().to_ascii_lowercase().as_str() {
76 "concept" => Ok(Self::Concept),
77 "document" | "doc" | "paper" => Ok(Self::Document),
78 "dataset" | "data" | "benchmark" => Ok(Self::Dataset),
79 "project" | "repo" | "crate" | "library" | "lib" => Ok(Self::Project),
80 "person" | "author" | "researcher" => Ok(Self::Person),
81 "org" | "organization" | "organisation" | "lab" | "company" => Ok(Self::Org),
82 other => Err(crate::error::UnknownVariant::new(
83 "entity_kind",
84 other,
85 ENTITY_KIND_VALID,
86 )),
87 }
88 }
89}
90
91#[derive(Clone, Debug)]
93#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
94pub struct Entity {
95 #[cfg_attr(feature = "serde", serde(flatten))]
96 pub header: Header,
97 pub kind: EntityKind,
98 pub name: String,
99 pub description: Option<String>,
100 pub properties: BTreeMap<String, PropertyValue>,
101 pub tags: Vec<String>,
102 pub deleted_at: Option<Timestamp>,
103}
104
105#[derive(Clone, Debug)]
107#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
108pub struct Link {
109 pub id: Id128,
110 pub source: Id128,
111 pub target: Id128,
112 pub relation: EdgeRelation,
113 pub properties: BTreeMap<String, PropertyValue>,
114 pub weight: f64,
115}
116
117#[derive(Clone, Debug, PartialEq)]
122#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
123#[cfg_attr(feature = "serde", serde(untagged))]
124pub enum PropertyValue {
125 String(String),
126 Integer(i64),
127 Float(f64),
128 Boolean(bool),
129 Array(Vec<PropertyValue>),
130 Object(BTreeMap<String, PropertyValue>),
131 Null,
132}
133
134impl fmt::Display for PropertyValue {
135 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
136 match self {
137 Self::String(s) => f.write_str(s),
138 Self::Integer(n) => write!(f, "{n}"),
139 Self::Float(n) => write!(f, "{n}"),
140 Self::Boolean(b) => write!(f, "{b}"),
141 Self::Array(arr) => write!(f, "[{} items]", arr.len()),
142 Self::Object(obj) => write!(f, "{{{} keys}}", obj.len()),
143 Self::Null => f.write_str("null"),
144 }
145 }
146}
147
148#[cfg(test)]
149mod tests {
150 use super::*;
151 use crate::{Namespace, Timestamp};
152
153 #[test]
154 fn entity_with_properties() {
155 let mut props = BTreeMap::new();
156 props.insert("role".into(), PropertyValue::String("engineer".into()));
157 props.insert("age".into(), PropertyValue::Integer(30));
158
159 let entity = Entity {
160 header: Header::new(
161 Id128::from_u128(1),
162 Namespace::default(),
163 Timestamp::from_secs(1700000000),
164 ),
165 kind: EntityKind::Person,
166 name: "Ocean".into(),
167 description: None,
168 properties: props,
169 tags: alloc::vec![],
170 deleted_at: None,
171 };
172 assert_eq!(entity.kind, EntityKind::Person);
173 assert_eq!(entity.kind.name(), "person");
174 assert_eq!(entity.properties.len(), 2);
175 }
176
177 #[test]
178 fn entity_kind_default_is_concept() {
179 assert_eq!(EntityKind::default(), EntityKind::Concept);
180 }
181
182 #[test]
183 fn entity_kind_display_roundtrip() {
184 for kind in EntityKind::ALL {
185 let s = alloc::format!("{kind}");
186 let parsed = EntityKind::from_str(&s).unwrap();
187 assert_eq!(parsed, kind);
188 }
189 }
190
191 #[test]
192 fn entity_kind_from_str_aliases() {
193 assert_eq!(EntityKind::from_str("doc").unwrap(), EntityKind::Document);
194 assert_eq!(EntityKind::from_str("paper").unwrap(), EntityKind::Document);
195 assert_eq!(
196 EntityKind::from_str("benchmark").unwrap(),
197 EntityKind::Dataset
198 );
199 assert_eq!(EntityKind::from_str("repo").unwrap(), EntityKind::Project);
200 assert_eq!(EntityKind::from_str("author").unwrap(), EntityKind::Person);
201 assert_eq!(EntityKind::from_str("lab").unwrap(), EntityKind::Org);
202 }
203
204 #[test]
205 fn entity_kind_from_str_case_insensitive() {
206 assert_eq!(
207 EntityKind::from_str("CONCEPT").unwrap(),
208 EntityKind::Concept
209 );
210 assert_eq!(EntityKind::from_str("Person").unwrap(), EntityKind::Person);
211 }
212
213 #[test]
214 fn entity_kind_from_str_unknown_errors() {
215 let err = EntityKind::from_str("gadget").unwrap_err();
216 assert_eq!(err.domain, "entity_kind");
217 assert_eq!(err.value, "gadget");
218 assert!(err.valid.contains(&"concept"));
219 }
220
221 #[test]
222 fn link_construction() {
223 let link = Link {
224 id: Id128::from_u128(100),
225 source: Id128::from_u128(1),
226 target: Id128::from_u128(2),
227 relation: EdgeRelation::Extends,
228 properties: BTreeMap::new(),
229 weight: 1.0,
230 };
231 assert_eq!(link.relation, EdgeRelation::Extends);
232 }
233}