use std::collections::BTreeMap;
use crate::error::{Error, Result};
use crate::key::Key;
use crate::keyword::FitsKeyword;
use crate::property::Property;
use crate::value::{FromField, IntoValue};
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct StructuralHints {
pub geometry: String,
pub sample_format: String,
pub color_space: String,
}
impl Default for StructuralHints {
fn default() -> Self {
Self {
geometry: "1:1:1".to_owned(),
sample_format: "UInt8".to_owned(),
color_space: "Gray".to_owned(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Header {
pub(crate) keywords: Vec<FitsKeyword>,
pub(crate) properties: BTreeMap<String, Property>,
}
impl Header {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn get<'a, T: FromField>(&self, key: impl Into<Key<'a>>) -> Result<Option<T>> {
Ok(self
.resolve(key.into())?
.and_then(|i| self.keywords[i].get::<T>()))
}
pub fn get_str<'a>(&self, key: impl Into<Key<'a>>) -> Result<Option<&str>> {
Ok(self
.resolve(key.into())?
.map(|i| self.keywords[i].value_str()))
}
pub fn get_f64<'a>(&self, key: impl Into<Key<'a>>) -> Result<Option<f64>> {
self.get(key)
}
pub fn get_i64<'a>(&self, key: impl Into<Key<'a>>) -> Result<Option<i64>> {
self.get(key)
}
pub fn get_u32<'a>(&self, key: impl Into<Key<'a>>) -> Result<Option<u32>> {
self.get(key)
}
pub fn get_bool<'a>(&self, key: impl Into<Key<'a>>) -> Result<Option<bool>> {
self.get(key)
}
pub fn get_datetime<'a>(
&self,
key: impl Into<Key<'a>>,
) -> Result<Option<time::PrimitiveDateTime>> {
self.get(key)
}
pub fn get_all<T: FromField>(&self, name: &str) -> Vec<T> {
self.indices(name)
.filter_map(|i| self.keywords[i].get::<T>())
.collect()
}
#[must_use]
pub fn count(&self, name: &str) -> usize {
self.indices(name).count()
}
#[must_use]
pub fn keywords(&self) -> &[FitsKeyword] {
&self.keywords
}
pub fn iter(&self) -> std::slice::Iter<'_, FitsKeyword> {
self.keywords.iter()
}
pub fn set<'a>(&mut self, key: impl Into<Key<'a>>, value: impl IntoValue) -> Result<()> {
let key = key.into();
let value = value.into_value();
match key {
Key::Name(name) => match self.resolve(Key::Name(name))? {
Some(i) => self.keywords[i].value = value,
None => {
Self::validate_name(name)?;
self.keywords.push(FitsKeyword {
name: name.to_owned(),
value,
comment: String::new(),
});
}
},
Key::Nth(name, n) => {
let i = self.require_nth(name, n)?;
self.keywords[i].value = value;
}
}
Ok(())
}
pub fn append(&mut self, name: &str, value: impl IntoValue) -> Result<()> {
Self::validate_name(name)?;
self.keywords.push(FitsKeyword {
name: name.to_owned(),
value: value.into_value(),
comment: String::new(),
});
Ok(())
}
pub fn set_comment<'a>(
&mut self,
key: impl Into<Key<'a>>,
comment: impl Into<String>,
) -> Result<bool> {
match self.resolve(key.into())? {
Some(i) => {
self.keywords[i].comment = comment.into();
Ok(true)
}
None => Ok(false),
}
}
pub fn set_with_comment<'a>(
&mut self,
key: impl Into<Key<'a>>,
value: impl IntoValue,
comment: impl Into<String>,
) -> Result<()> {
let key = key.into();
self.set(key, value)?;
if let Some(i) = self.resolve(key)? {
self.keywords[i].comment = comment.into();
}
Ok(())
}
pub fn remove<'a>(&mut self, key: impl Into<Key<'a>>) -> Result<bool> {
match self.resolve(key.into())? {
Some(i) => {
self.keywords.remove(i);
Ok(true)
}
None => Ok(false),
}
}
pub fn remove_all(&mut self, name: &str) -> usize {
let before = self.keywords.len();
self.keywords.retain(|k| !k.name.eq_ignore_ascii_case(name));
before - self.keywords.len()
}
pub fn set_many<'a, V, I>(&mut self, entries: I) -> Result<()>
where
V: IntoValue,
I: IntoIterator<Item = (&'a str, V)>,
{
let entries: Vec<(&str, V)> = entries.into_iter().collect();
for (name, _) in &entries {
Self::validate_name(name)?;
let count = self.count(name);
if count > 1 {
return Err(Error::Ambiguous {
name: (*name).to_owned(),
count,
});
}
}
for (name, value) in entries {
match self.first_index(name) {
Some(i) => self.keywords[i].value = value.into_value(),
None => self.keywords.push(FitsKeyword {
name: name.to_owned(),
value: value.into_value(),
comment: String::new(),
}),
}
}
Ok(())
}
pub fn remove_many<'a, I: IntoIterator<Item = &'a str>>(&mut self, names: I) -> Result<usize> {
let names: Vec<&str> = names.into_iter().collect();
for name in &names {
let count = self.count(name);
if count > 1 {
return Err(Error::Ambiguous {
name: (*name).to_owned(),
count,
});
}
}
let mut removed = 0;
for name in names {
if let Some(i) = self.first_index(name) {
self.keywords.remove(i);
removed += 1;
}
}
Ok(removed)
}
#[must_use]
pub fn properties(&self) -> &BTreeMap<String, Property> {
&self.properties
}
#[must_use]
pub fn property(&self, id: &str) -> Option<&str> {
self.properties.get(id).map(|p| p.value.as_str())
}
#[must_use]
pub fn property_get<T: FromField>(&self, id: &str) -> Option<T> {
self.properties
.get(id)
.and_then(|p| T::from_field(&p.value))
}
pub fn set_property(&mut self, id: impl Into<String>, value: impl Into<String>) -> Result<()> {
let id = id.into();
Self::validate_property_id(&id)?;
self.properties.entry(id).or_default().value = value.into();
Ok(())
}
pub fn set_property_with_type(
&mut self,
id: impl Into<String>,
value: impl Into<String>,
type_: impl Into<String>,
) -> Result<()> {
let id = id.into();
Self::validate_property_id(&id)?;
let p = self.properties.entry(id).or_default();
p.value = value.into();
p.type_ = type_.into();
Ok(())
}
pub fn remove_property(&mut self, id: &str) -> bool {
self.properties.remove(id).is_some()
}
fn indices<'s>(&'s self, name: &'s str) -> impl Iterator<Item = usize> + 's {
self.keywords
.iter()
.enumerate()
.filter(move |(_, k)| k.name.eq_ignore_ascii_case(name))
.map(|(i, _)| i)
}
fn first_index(&self, name: &str) -> Option<usize> {
self.indices(name).next()
}
fn resolve(&self, key: Key) -> Result<Option<usize>> {
match key {
Key::Name(name) => {
let mut it = self.indices(name);
let first = it.next();
if first.is_some() && it.next().is_some() {
return Err(Error::Ambiguous {
name: name.to_owned(),
count: self.count(name),
});
}
Ok(first)
}
Key::Nth(name, n) => {
let indices: Vec<usize> = self.indices(name).collect();
match indices.get(n) {
Some(&i) => Ok(Some(i)),
None if indices.is_empty() => Ok(None),
None => Err(Error::IndexOutOfRange {
name: name.to_owned(),
index: n,
count: indices.len(),
}),
}
}
}
}
fn require_nth(&self, name: &str, n: usize) -> Result<usize> {
self.resolve(Key::Nth(name, n))?
.ok_or_else(|| Error::IndexOutOfRange {
name: name.to_owned(),
index: n,
count: 0,
})
}
fn validate_name(name: &str) -> Result<()> {
if name.is_empty() {
return Err(Error::InvalidName {
name: name.to_owned(),
reason: "empty",
});
}
if name.len() > 8 {
return Err(Error::InvalidName {
name: name.to_owned(),
reason: "exceeds 8 characters",
});
}
if !name
.bytes()
.all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_')
{
return Err(Error::InvalidName {
name: name.to_owned(),
reason: "must be ASCII letters, digits, `-`, or `_`",
});
}
Ok(())
}
fn validate_property_id(id: &str) -> Result<()> {
if id.is_empty() {
return Err(Error::InvalidName {
name: id.to_owned(),
reason: "empty",
});
}
if !id
.bytes()
.all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b':')
{
return Err(Error::InvalidName {
name: id.to_owned(),
reason: "property id must be ASCII alphanumeric, `_`, or `:`",
});
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn validate_name_rules() {
assert!(Header::validate_name("GAIN").is_ok());
assert!(Header::validate_name("DATE-OBS").is_ok());
assert!(Header::validate_name("lower_k").is_ok());
assert!(Header::validate_name("EIGHTCHR").is_ok());
assert!(Header::validate_name("").is_err());
assert!(Header::validate_name("NINECHARS").is_err());
assert!(Header::validate_name("BAD KEY").is_err());
assert!(Header::validate_name("NAME!").is_err());
}
#[test]
fn validate_property_id_rules() {
assert!(Header::validate_property_id("Instrument:Telescope:FocalLength").is_ok());
assert!(Header::validate_property_id("A_b:9").is_ok());
assert!(Header::validate_property_id("").is_err());
assert!(Header::validate_property_id("bad id!").is_err());
assert!(Header::validate_property_id("hy-phen").is_err());
}
#[test]
fn nth_write_on_absent_name_errors() {
let mut h = Header::new();
assert!(matches!(
h.set(("MISSING", 0), 1_i64),
Err(Error::IndexOutOfRange { count: 0, .. })
));
}
#[test]
fn set_with_comment_creates_and_updates() {
let mut h = Header::new();
h.set_with_comment("GAIN", 100_i64, "sensor gain").unwrap();
assert_eq!(h.get_i64("GAIN").unwrap(), Some(100));
assert_eq!(h.keywords()[0].comment, "sensor gain");
h.set_with_comment("GAIN", 200_i64, "updated").unwrap();
assert_eq!(h.get_i64("GAIN").unwrap(), Some(200));
assert_eq!(h.keywords()[0].comment, "updated");
h.append("HISTORY", "a").unwrap();
h.append("HISTORY", "b").unwrap();
assert!(matches!(
h.set_with_comment("HISTORY", "x", "c"),
Err(Error::Ambiguous { .. })
));
}
#[test]
fn set_comment_on_absent_keyword_reports_not_found() {
let mut h = Header::new();
assert!(!h.set_comment("MISSING", "c").unwrap());
}
#[test]
fn remove_all_clears_every_occurrence() {
let mut h = Header::new();
h.append("HISTORY", "a").unwrap();
h.append("HISTORY", "b").unwrap();
h.set("GAIN", 1_i64).unwrap();
assert_eq!(h.remove_all("history"), 2); assert_eq!(h.count("HISTORY"), 0);
assert_eq!(h.remove_all("HISTORY"), 0);
assert_eq!(h.get_i64("GAIN").unwrap(), Some(1));
}
#[test]
fn iter_preserves_document_order() {
let mut h = Header::new();
h.set("A", 1_i64).unwrap();
h.set("B", 2_i64).unwrap();
h.set("C", 3_i64).unwrap();
let names: Vec<&str> = h.iter().map(|k| k.name.as_str()).collect();
assert_eq!(names, ["A", "B", "C"]);
}
#[test]
fn string_keys_are_accepted() {
let mut h = Header::new();
let key = String::from("GAIN");
h.set(&key, 100_i64).unwrap();
assert_eq!(h.get_i64(&key).unwrap(), Some(100));
}
#[test]
fn generic_get_reads_string() {
let mut h = Header::new();
h.set("OBJECT", "M31").unwrap();
assert_eq!(h.get::<String>("OBJECT").unwrap(), Some("M31".to_owned()));
}
}