1use crate::error::FitsError;
4use crate::key::Key;
5use crate::record::{is_commentary_keyword, validate_keyword, validate_keyword_raw, Record, Value};
6use crate::value::{FromCard, IntoValue};
7use crate::write::{self, StructuralHints};
8
9#[derive(Clone, Debug, Default, PartialEq)]
13#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
14pub struct Header {
15 records: Vec<Record>,
16}
17
18impl Header {
19 pub fn new() -> Self {
21 Header::default()
22 }
23
24 pub(crate) fn from_records(records: Vec<Record>) -> Self {
26 Header { records }
27 }
28
29 pub fn cards(&self) -> &[Record] {
31 &self.records
32 }
33
34 pub fn iter(&self) -> impl Iterator<Item = &Record> {
36 self.records.iter()
37 }
38
39 pub fn count(&self, name: &str) -> usize {
41 self.records
42 .iter()
43 .filter(|r| r.keyword() == Some(name))
44 .count()
45 }
46
47 fn resolve(&self, key: &Key) -> Result<Option<usize>, FitsError> {
49 let name = key.name();
50 let indices: Vec<usize> = self
51 .records
52 .iter()
53 .enumerate()
54 .filter(|(_, r)| r.keyword() == Some(name))
55 .map(|(i, _)| i)
56 .collect();
57 match key.occurrence() {
58 Some(n) => Ok(indices.get(n).copied()),
59 None => match indices.len() {
60 0 => Ok(None),
61 1 => Ok(Some(indices[0])),
62 count => Err(FitsError::AmbiguousKeyword {
63 keyword: name.to_string(),
64 count,
65 }),
66 },
67 }
68 }
69
70 pub fn get<T: FromCard>(&self, key: impl Into<Key>) -> Result<Option<T>, FitsError> {
83 Ok(self
84 .resolve(&key.into())?
85 .and_then(|i| T::from_card(&self.records[i])))
86 }
87
88 pub fn get_str(&self, key: impl Into<Key>) -> Result<Option<&str>, FitsError> {
90 Ok(self
91 .resolve(&key.into())?
92 .and_then(|i| self.records[i].str_content()))
93 }
94
95 pub fn get_all<T: FromCard>(&self, name: &str) -> Vec<T> {
97 self.records
98 .iter()
99 .filter(|r| r.keyword() == Some(name))
100 .filter_map(T::from_card)
101 .collect()
102 }
103
104 fn make_record(name: &str, value: Value) -> Record {
105 if is_commentary_keyword(name) {
106 let text = match value {
107 Value::Str(s) | Value::Literal(s) => s,
108 };
109 Record::commentary(name, text)
110 } else {
111 Record::value(name, value, None)
112 }
113 }
114
115 fn set_inner(&mut self, key: Key, value: Value, raw: bool) -> Result<(), FitsError> {
116 let name = key.name().to_string();
117 if raw {
118 validate_keyword_raw(&name)?;
119 } else {
120 validate_keyword(&name)?;
121 }
122 match self.resolve(&key)? {
123 Some(i) => {
124 self.records[i].replace_value(value);
125 Ok(())
126 }
127 None => match key.occurrence() {
128 Some(n) => Err(FitsError::OccurrenceOutOfRange {
129 keyword: name.clone(),
130 occurrence: n,
131 count: self.count(&name),
132 }),
133 None => {
134 self.records.push(Self::make_record(&name, value));
135 Ok(())
136 }
137 },
138 }
139 }
140
141 pub fn set(&mut self, key: impl Into<Key>, value: impl IntoValue) -> Result<(), FitsError> {
158 self.set_inner(key.into(), value.into_value(), false)
159 }
160
161 pub fn set_raw(&mut self, keyword: &str, value: impl IntoValue) -> Result<(), FitsError> {
163 self.set_inner(Key::Name(keyword.to_string()), value.into_value(), true)
164 }
165
166 pub fn append(&mut self, name: &str, value: impl IntoValue) -> Result<(), FitsError> {
178 validate_keyword(name)?;
179 self.records
180 .push(Self::make_record(name, value.into_value()));
181 Ok(())
182 }
183
184 pub fn set_comment(
187 &mut self,
188 key: impl Into<Key>,
189 comment: impl Into<String>,
190 ) -> Result<(), FitsError> {
191 if let Some(i) = self.resolve(&key.into())? {
192 self.records[i].set_comment(Some(comment.into()));
193 }
194 Ok(())
195 }
196
197 pub fn remove(&mut self, key: impl Into<Key>) -> Result<bool, FitsError> {
209 match self.resolve(&key.into())? {
210 Some(i) => {
211 self.records.remove(i);
212 Ok(true)
213 }
214 None => Ok(false),
215 }
216 }
217
218 pub fn set_many<K, V>(
232 &mut self,
233 entries: impl IntoIterator<Item = (K, V)>,
234 ) -> Result<(), FitsError>
235 where
236 K: Into<Key>,
237 V: IntoValue,
238 {
239 let items: Vec<(Key, Value)> = entries
240 .into_iter()
241 .map(|(k, v)| (k.into(), v.into_value()))
242 .collect();
243 for (k, _) in &items {
245 validate_keyword(k.name())?;
246 if let Some(n) = k.occurrence() {
247 if self.resolve(k)?.is_none() {
248 return Err(FitsError::OccurrenceOutOfRange {
249 keyword: k.name().to_string(),
250 occurrence: n,
251 count: self.count(k.name()),
252 });
253 }
254 } else {
255 self.resolve(k)?;
257 }
258 }
259 for (k, v) in items {
260 self.set_inner(k, v, false)?;
261 }
262 Ok(())
263 }
264
265 pub fn remove_many<K: Into<Key>>(
267 &mut self,
268 keys: impl IntoIterator<Item = K>,
269 ) -> Result<usize, FitsError> {
270 let keys: Vec<Key> = keys.into_iter().map(Into::into).collect();
271 for k in &keys {
272 self.resolve(k)?;
273 }
274 let mut removed = 0;
275 for k in keys {
276 if self.remove(k)? {
277 removed += 1;
278 }
279 }
280 Ok(removed)
281 }
282
283 pub fn to_header_bytes(&self) -> Vec<u8> {
296 write::to_header_bytes(self)
297 }
298
299 pub fn to_bytes(&self, structural: &StructuralHints) -> Result<Vec<u8>, FitsError> {
316 write::to_bytes(self, structural)
317 }
318}
319
320#[cfg(test)]
321mod tests {
322 use super::*;
323
324 #[test]
325 fn set_routes_commentary_keywords_to_commentary_records() {
326 let mut h = Header::new();
327 h.set("COMMENT", "a note").unwrap();
328 h.set("HISTORY", "step 1").unwrap();
329 assert!(matches!(
330 h.cards()[0].kind,
331 crate::record::RecordKind::Commentary { .. }
332 ));
333 assert_eq!(h.get_all::<String>("HISTORY"), vec!["step 1".to_string()]);
334 }
335
336 #[test]
337 fn get_all_skips_unconvertible_values() {
338 let mut h = Header::new();
339 h.append("GAIN", 100).unwrap();
340 h.append("GAIN", "not a number").unwrap();
341 h.append("GAIN", 200).unwrap();
342 assert_eq!(h.get_all::<i64>("GAIN"), vec![100, 200]);
343 assert_eq!(h.count("GAIN"), 3);
344 }
345
346 #[test]
347 fn set_comment_on_absent_key_is_noop() {
348 let mut h = Header::new();
349 h.set_comment("NOPE", "x").unwrap();
350 assert!(h.cards().is_empty());
351 }
352
353 #[test]
354 fn remove_returns_false_when_absent() {
355 let mut h = Header::new();
356 assert!(!h.remove("NOPE").unwrap());
357 }
358
359 #[test]
360 fn remove_many_aborts_on_ambiguity_before_removing() {
361 let mut h = Header::new();
362 h.set("A", 1).unwrap();
363 h.append("DUP", 1).unwrap();
364 h.append("DUP", 2).unwrap();
365 let before = h.clone();
366 assert!(matches!(
367 h.remove_many(["A", "DUP"]),
368 Err(FitsError::AmbiguousKeyword { .. })
369 ));
370 assert_eq!(h, before, "nothing may be removed on a rejected batch");
371
372 assert_eq!(h.remove_many(["A", "MISSING"]).unwrap(), 1);
373 }
374
375 #[test]
376 fn set_many_accepts_occurrence_keys() {
377 let mut h = Header::new();
378 h.append("GAIN", 1).unwrap();
379 h.append("GAIN", 2).unwrap();
380 h.set_many([(("GAIN", 0), 10), (("GAIN", 1), 20)]).unwrap();
381 assert_eq!(h.get_all::<i64>("GAIN"), vec![10, 20]);
382 }
383
384 #[test]
385 fn iter_matches_cards() {
386 let mut h = Header::new();
387 h.set("A", 1).unwrap();
388 h.set("B", 2).unwrap();
389 assert_eq!(h.iter().count(), 2);
390 let names: Vec<_> = h.iter().filter_map(|r| r.keyword()).collect();
391 assert_eq!(names, vec!["A", "B"]);
392 }
393
394 #[test]
395 fn get_on_missing_key_is_ok_none() {
396 let h = Header::new();
397 assert_eq!(h.get::<i64>("NOPE").unwrap(), None);
398 assert_eq!(h.get_str("NOPE").unwrap(), None);
399 assert_eq!(h.get::<i64>(("NOPE", 3)).unwrap(), None);
400 }
401}