onnx_extractor/
attribute_value.rs

1use prost::bytes::Bytes;
2
3use crate::{Error, Graph, Tensor};
4
5/// ONNX attribute values
6#[derive(Debug)]
7pub enum AttributeValue {
8    Float(f32),
9    Int(i64),
10    String(Bytes),
11    Tensor(Box<Tensor>),
12    Graph(Box<Graph>),
13    Floats(Vec<f32>),
14    Ints(Vec<i64>),
15    Strings(Vec<Bytes>),
16    Tensors(Box<[Tensor]>),
17    Graphs(Box<[Graph]>),
18}
19
20impl AttributeValue {
21    /// Get string value as raw bytes without UTF-8 validation.
22    pub fn as_string(&self) -> Option<&Bytes> {
23        match self {
24            AttributeValue::String(s) => Some(s),
25            _ => None,
26        }
27    }
28
29    /// Get string value as a validated UTF-8 `&str`.
30    ///
31    /// Returns `Err` if the variant is not `String` or if the bytes are not valid UTF-8.
32    /// The returned `&str` borrows directly from the underlying buffer with no copy.
33    pub fn as_string_validated(&self) -> Result<&str, Error> {
34        match self {
35            AttributeValue::String(s) => Ok(std::str::from_utf8(s)?),
36            _ => Err(Error::MissingField("string attribute")),
37        }
38    }
39
40    /// Extract string value as owned `Bytes`.
41    pub fn into_string(self) -> Option<Bytes> {
42        match self {
43            AttributeValue::String(s) => Some(s),
44            _ => None,
45        }
46    }
47
48    /// Get string array value as raw bytes without UTF-8 validation.
49    pub fn as_strings(&self) -> Option<&[Bytes]> {
50        match self {
51            AttributeValue::Strings(s) => Some(s),
52            _ => None,
53        }
54    }
55
56    /// Get string array value as validated `&str` entries.
57    ///
58    /// Returns `Err` if the variant is not `Strings` or if any entry is not valid UTF-8.
59    /// Each `&str` borrows directly from the underlying buffer with no copy,
60    /// but the returned `Box<[&str]>` is a new allocation for the pointer array.
61    pub fn as_strings_validated(&self) -> Result<Box<[&str]>, Error> {
62        match self {
63            AttributeValue::Strings(strings) => strings
64                .iter()
65                .map(|s| Ok(std::str::from_utf8(s)?))
66                .collect(),
67            _ => Err(Error::MissingField("strings attribute")),
68        }
69    }
70
71    /// Extract string array value as owned `Vec<Bytes>`.
72    pub fn into_strings(self) -> Option<Vec<Bytes>> {
73        match self {
74            AttributeValue::Strings(s) => Some(s),
75            _ => None,
76        }
77    }
78
79    /// Get float value if the variant is `Float`.
80    pub fn as_float(&self) -> Option<f32> {
81        match self {
82            AttributeValue::Float(f) => Some(*f),
83            _ => None,
84        }
85    }
86
87    /// Get integer value if the variant is `Int`.
88    pub fn as_int(&self) -> Option<i64> {
89        match self {
90            AttributeValue::Int(i) => Some(*i),
91            _ => None,
92        }
93    }
94
95    /// Borrow float slice if the variant is `Floats`.
96    pub fn as_floats(&self) -> Option<&[f32]> {
97        match self {
98            AttributeValue::Floats(f) => Some(f),
99            _ => None,
100        }
101    }
102
103    /// Extract float vector as owned `Vec<f32>` if the variant is `Floats`.
104    pub fn into_floats(self) -> Option<Vec<f32>> {
105        match self {
106            AttributeValue::Floats(f) => Some(f),
107            _ => None,
108        }
109    }
110
111    /// Borrow integer slice if the variant is `Ints`.
112    pub fn as_ints(&self) -> Option<&[i64]> {
113        match self {
114            AttributeValue::Ints(i) => Some(i),
115            _ => None,
116        }
117    }
118
119    /// Extract integer vector as owned `Vec<i64>` if the variant is `Ints`.
120    pub fn into_ints(self) -> Option<Vec<i64>> {
121        match self {
122            AttributeValue::Ints(i) => Some(i),
123            _ => None,
124        }
125    }
126
127    /// Borrow tensor if the variant is `Tensor`.
128    pub fn as_tensor(&self) -> Option<&Tensor> {
129        match self {
130            AttributeValue::Tensor(t) => Some(t),
131            _ => None,
132        }
133    }
134
135    /// Extract tensor as owned `Tensor` if the variant is `Tensor`.
136    pub fn into_tensor(self) -> Option<Tensor> {
137        match self {
138            AttributeValue::Tensor(t) => Some(*t),
139            _ => None,
140        }
141    }
142
143    /// Borrow subgraph if the variant is `Graph`.
144    pub fn as_graph(&self) -> Option<&Graph> {
145        match self {
146            AttributeValue::Graph(g) => Some(g),
147            _ => None,
148        }
149    }
150
151    /// Extract subgraph as owned `Graph` if the variant is `Graph`.
152    pub fn into_graph(self) -> Option<Graph> {
153        match self {
154            AttributeValue::Graph(g) => Some(*g),
155            _ => None,
156        }
157    }
158
159    /// Borrow tensor slice if the variant is `Tensors`.
160    pub fn as_tensors(&self) -> Option<&[Tensor]> {
161        match self {
162            AttributeValue::Tensors(t) => Some(t),
163            _ => None,
164        }
165    }
166
167    /// Extract tensor slice as owned `Box<[Tensor]>` if the variant is `Tensors`.
168    pub fn into_tensors(self) -> Option<Box<[Tensor]>> {
169        match self {
170            AttributeValue::Tensors(t) => Some(t),
171            _ => None,
172        }
173    }
174
175    /// Borrow subgraph slice if the variant is `Graphs`.
176    pub fn as_graphs(&self) -> Option<&[Graph]> {
177        match self {
178            AttributeValue::Graphs(g) => Some(g),
179            _ => None,
180        }
181    }
182
183    /// Extract subgraph slice as owned `Box<[Graph]>` if the variant is `Graphs`.
184    pub fn into_graphs(self) -> Option<Box<[Graph]>> {
185        match self {
186            AttributeValue::Graphs(g) => Some(g),
187            _ => None,
188        }
189    }
190}
onnx_extractor/attribute_value.rs

onnx_extractor/
attribute_value.rs
