Skip to main content

nyx_space/io/
mod.rs

1/*
2    Nyx, blazing fast astrodynamics
3    Copyright (C) 2018-onwards Christopher Rabotin <christopher.rabotin@gmail.com>
4
5    This program is free software: you can redistribute it and/or modify
6    it under the terms of the GNU Affero General Public License as published
7    by the Free Software Foundation, either version 3 of the License, or
8    (at your option) any later version.
9
10    This program is distributed in the hope that it will be useful,
11    but WITHOUT ANY WARRANTY; without even the implied warranty of
12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13    GNU Affero General Public License for more details.
14
15    You should have received a copy of the GNU Affero General Public License
16    along with this program.  If not, see <https://www.gnu.org/licenses/>.
17*/
18
19use crate::md::StateParameter;
20use crate::time::Epoch;
21use arrow::error::ArrowError;
22use log::debug;
23use parquet::errors::ParquetError;
24use snafu::prelude::*;
25pub(crate) mod watermark;
26use hifitime::Duration;
27use hifitime::prelude::{Format, Formatter};
28use serde::de::DeserializeOwned;
29use serde::{Deserialize, Deserializer};
30use serde::{Serialize, Serializer};
31use std::collections::{BTreeMap, HashMap};
32use std::fmt::Debug;
33use std::fs::File;
34use std::io::BufReader;
35use std::path::{Path, PathBuf};
36use std::str::FromStr;
37use typed_builder::TypedBuilder;
38
39pub mod gravity;
40pub mod space_weather;
41
42use std::io;
43
44#[cfg(feature = "python")]
45use pyo3::prelude::*;
46#[cfg(feature = "python")]
47mod python;
48
49/// Configuration for exporting from Nyx to local disk.
50///
51/// :type timestamped: bool
52#[derive(Clone, Debug, Default, Serialize, Deserialize, TypedBuilder, PartialEq)]
53#[builder(doc)]
54#[cfg_attr(feature = "python", pyclass(from_py_object, eq))]
55pub struct ExportCfg {
56    /// Fields to export, if unset, defaults to all possible fields.
57    #[builder(default, setter(strip_option))]
58    pub fields: Option<Vec<StateParameter>>,
59    /// Start epoch to export, defaults to the start of the trajectory
60    #[builder(default, setter(strip_option))]
61    pub start_epoch: Option<Epoch>,
62    /// End epoch to export, defaults to the end of the trajectory
63    #[builder(default, setter(strip_option))]
64    pub end_epoch: Option<Epoch>,
65    /// An optional step, defaults to every state in the trajectory (which likely isn't equidistant)
66    #[builder(default, setter(strip_option))]
67    pub step: Option<Duration>,
68    /// Additional metadata to store in the Parquet metadata
69    #[builder(default, setter(strip_option))]
70    pub metadata: Option<HashMap<String, String>>,
71    /// Set to true to append the timestamp to the filename
72    #[builder(default)]
73    pub timestamp: bool,
74}
75
76impl ExportCfg {
77    /// Initialize a new configuration with the given metadata entries.
78    pub fn from_metadata(metadata: Vec<(String, String)>) -> Self {
79        let mut me = ExportCfg {
80            metadata: Some(HashMap::new()),
81            ..Default::default()
82        };
83        for (k, v) in metadata {
84            me.metadata.as_mut().unwrap().insert(k, v);
85        }
86        me
87    }
88
89    /// Initialize a new default configuration but timestamp the filename.
90    pub fn timestamped() -> Self {
91        Self {
92            timestamp: true,
93            ..Default::default()
94        }
95    }
96
97    pub fn append_field(&mut self, field: StateParameter) {
98        if let Some(fields) = self.fields.as_mut() {
99            fields.push(field);
100        } else {
101            self.fields = Some(vec![field]);
102        }
103    }
104
105    /// Modifies the provided path to include the timestamp if required.
106    pub(crate) fn actual_path<P: AsRef<Path>>(&self, path: P) -> PathBuf {
107        let mut path_buf = path.as_ref().to_path_buf();
108        if self.timestamp
109            && let Some(file_name) = path_buf.file_name()
110            && let Some(file_name_str) = file_name.to_str()
111            && let Some(extension) = path_buf.extension()
112        {
113            let stamp = Formatter::new(
114                Epoch::now().unwrap(),
115                Format::from_str("%Y-%m-%dT%H-%M-%S").unwrap(),
116            );
117            let ext = extension.to_str().unwrap();
118            let file_name = file_name_str.replace(&format!(".{ext}"), "");
119            let new_file_name = format!("{file_name}-{stamp}.{ext}");
120            path_buf.set_file_name(new_file_name);
121        };
122        path_buf
123    }
124}
125
126#[derive(Debug, Snafu)]
127#[snafu(visibility(pub(crate)))]
128pub enum ConfigError {
129    #[snafu(display("failed to read configuration file: {source}"))]
130    ReadError { source: io::Error },
131
132    #[snafu(display("failed to parse YAML configuration file: {source}"))]
133    ParseError { source: serde_yml::Error },
134
135    #[snafu(display("of invalid configuration: {msg}"))]
136    InvalidConfig { msg: String },
137}
138
139impl PartialEq for ConfigError {
140    /// No two configuration errors match
141    fn eq(&self, _other: &Self) -> bool {
142        false
143    }
144}
145
146#[derive(Debug, Snafu)]
147#[snafu(visibility(pub(crate)))]
148pub enum InputOutputError {
149    #[snafu(display("{action} encountered i/o error: {source}"))]
150    StdIOError {
151        source: io::Error,
152        action: &'static str,
153    },
154    #[snafu(display("missing required data {which}"))]
155    MissingData { which: String },
156    #[snafu(display("unknown data `{which}`"))]
157    UnsupportedData { which: String },
158    #[snafu(display("{action} encountered a Parquet error: {source}"))]
159    ParquetError {
160        source: ParquetError,
161        action: &'static str,
162    },
163    #[snafu(display("inconsistency detected: {msg}"))]
164    Inconsistency { msg: String },
165    #[snafu(display("{action} encountered an Arrow error: {source}"))]
166    ArrowError {
167        source: ArrowError,
168        action: &'static str,
169    },
170    #[snafu(display("error parsing `{data}` as Dhall config: {err}"))]
171    ParseDhall { data: String, err: String },
172    #[snafu(display("error serializing {what} to Dhall: {err}"))]
173    SerializeDhall { what: String, err: String },
174    #[snafu(display("empty dataset error when (de)serializing {action}"))]
175    EmptyDataset { action: &'static str },
176    #[snafu(display("CSV reading errors when {action}: {source}"))]
177    CsvData {
178        source: csv::Error,
179        action: &'static str,
180    },
181}
182
183impl PartialEq for InputOutputError {
184    fn eq(&self, _other: &Self) -> bool {
185        false
186    }
187}
188
189pub trait ConfigRepr: Debug + Sized + Serialize + DeserializeOwned {
190    /// Builds the configuration representation from the path to a yaml
191    fn load<P>(path: P) -> Result<Self, ConfigError>
192    where
193        P: AsRef<Path>,
194    {
195        let file = File::open(path).context(ReadSnafu)?;
196        let reader = BufReader::new(file);
197
198        serde_yml::from_reader(reader).context(ParseSnafu)
199    }
200
201    /// Builds a sequence of "Selves" from the provided path to a yaml
202    fn load_many<P>(path: P) -> Result<Vec<Self>, ConfigError>
203    where
204        P: AsRef<Path>,
205    {
206        let file = File::open(path).context(ReadSnafu)?;
207        let reader = BufReader::new(file);
208
209        serde_yml::from_reader(reader).context(ParseSnafu)
210    }
211
212    /// Builds a map of names to "selves" from the provided path to a yaml
213    fn load_named<P>(path: P) -> Result<BTreeMap<String, Self>, ConfigError>
214    where
215        P: AsRef<Path>,
216    {
217        let file = File::open(path).context(ReadSnafu)?;
218        let reader = BufReader::new(file);
219
220        serde_yml::from_reader(reader).context(ParseSnafu)
221    }
222
223    /// Builds a sequence of "Selves" from the provided string of a yaml
224    fn loads_many(data: &str) -> Result<Vec<Self>, ConfigError> {
225        debug!("Loading YAML:\n{data}");
226        serde_yml::from_str(data).context(ParseSnafu)
227    }
228
229    /// Builds a sequence of "Selves" from the provided string of a yaml
230    fn loads_named(data: &str) -> Result<BTreeMap<String, Self>, ConfigError> {
231        debug!("Loading YAML:\n{data}");
232        serde_yml::from_str(data).context(ParseSnafu)
233    }
234}
235
236pub(crate) fn epoch_to_str<S>(epoch: &Epoch, serializer: S) -> Result<S::Ok, S::Error>
237where
238    S: Serializer,
239{
240    serializer.serialize_str(&format!("{epoch}"))
241}
242
243/// A deserializer from Epoch string
244pub(crate) fn epoch_from_str<'de, D>(deserializer: D) -> Result<Epoch, D::Error>
245where
246    D: Deserializer<'de>,
247{
248    // implementation of the custom deserialization function
249    let s = String::deserialize(deserializer)?;
250    Epoch::from_str(&s).map_err(serde::de::Error::custom)
251}
252
253pub(crate) fn duration_to_str<S>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error>
254where
255    S: Serializer,
256{
257    serializer.serialize_str(&format!("{duration}"))
258}
259
260/// A deserializer from Duration string
261pub(crate) fn duration_from_str<'de, D>(deserializer: D) -> Result<Duration, D::Error>
262where
263    D: Deserializer<'de>,
264{
265    // implementation of the custom deserialization function
266    let s = String::deserialize(deserializer)?;
267    Duration::from_str(&s).map_err(serde::de::Error::custom)
268}
269
270pub(crate) fn maybe_duration_to_str<S>(
271    duration: &Option<Duration>,
272    serializer: S,
273) -> Result<S::Ok, S::Error>
274where
275    S: Serializer,
276{
277    if let Some(duration) = duration {
278        duration_to_str(duration, serializer)
279    } else {
280        serializer.serialize_none()
281    }
282}
283
284pub(crate) fn maybe_duration_from_str<'de, D>(deserializer: D) -> Result<Option<Duration>, D::Error>
285where
286    D: Deserializer<'de>,
287{
288    if let Ok(s) = String::deserialize(deserializer) {
289        if let Ok(duration) = Duration::from_str(&s) {
290            Ok(Some(duration))
291        } else {
292            Ok(None)
293        }
294    } else {
295        Ok(None)
296    }
297}