ixa/
report.rs

1use crate::context::Context;
2use crate::error::IxaError;
3use crate::people::ContextPeopleExt;
4use crate::{define_data_plugin, Tabulator};
5use crate::{error, trace};
6use crate::{HashMap, HashMapExt, PluginContext};
7use csv::Writer;
8use serde::Serializer;
9use std::any::TypeId;
10use std::cell::{RefCell, RefMut};
11use std::env;
12use std::fs::File;
13use std::path::PathBuf;
14
15// * file_prefix: precedes the report name in the filename. An example of a
16// potential prefix might be scenario or simulation name
17// * directory: location that the CSVs are written to. An example of this might
18// be /data/
19// * overwrite: if true, will overwrite existing files in the same location
20pub struct ConfigReportOptions {
21    pub file_prefix: String,
22    pub output_dir: PathBuf,
23    pub overwrite: bool,
24}
25
26impl ConfigReportOptions {
27    #[must_use]
28    #[allow(clippy::missing_panics_doc)]
29    pub fn new() -> Self {
30        trace!("new ConfigReportOptions");
31        // Sets the defaults
32        ConfigReportOptions {
33            file_prefix: String::new(),
34            output_dir: env::current_dir().unwrap(),
35            overwrite: false,
36        }
37    }
38    /// Sets the file prefix option (e.g., "report_")
39    pub fn file_prefix(&mut self, file_prefix: String) -> &mut ConfigReportOptions {
40        trace!("setting report prefix to {file_prefix}");
41        self.file_prefix = file_prefix;
42        self
43    }
44    /// Sets the directory where reports will be output
45    pub fn directory(&mut self, directory: PathBuf) -> &mut ConfigReportOptions {
46        trace!("setting report directory to {directory:?}");
47        self.output_dir = directory;
48        self
49    }
50    /// Sets whether to overwrite existing reports of the same name if they exist
51    pub fn overwrite(&mut self, overwrite: bool) -> &mut ConfigReportOptions {
52        trace!("setting report overwrite {overwrite}");
53        self.overwrite = overwrite;
54        self
55    }
56}
57
58impl Default for ConfigReportOptions {
59    fn default() -> Self {
60        Self::new()
61    }
62}
63
64pub trait Report: 'static {
65    // Returns report type
66    fn type_id(&self) -> TypeId;
67    // Serializes the data with the correct writer
68    fn serialize(&self, writer: &mut Writer<File>);
69}
70
71/// Use this macro to define a unique report type
72#[macro_export]
73macro_rules! define_report {
74    ($name:ident) => {
75        impl $crate::Report for $name {
76            fn type_id(&self) -> std::any::TypeId {
77                std::any::TypeId::of::<$name>()
78            }
79
80            fn serialize(&self, writer: &mut $crate::csv::Writer<std::fs::File>) {
81                writer.serialize(self).unwrap();
82            }
83        }
84    };
85}
86
87/// # Errors
88/// function will return Error if it fails to `serialize_str`
89#[allow(clippy::trivially_copy_pass_by_ref)]
90#[allow(dead_code)]
91pub fn serialize_f64<S, const N: usize>(value: &f64, serializer: S) -> Result<S::Ok, S::Error>
92where
93    S: Serializer,
94{
95    let formatted = format!("{value:.N$}");
96    serializer.serialize_str(&formatted)
97}
98
99/// # Errors
100/// function will return Error if it fails to `serialize_str`
101#[allow(clippy::trivially_copy_pass_by_ref)]
102#[allow(dead_code)]
103pub fn serialize_f32<S, const N: usize>(value: &f32, serializer: S) -> Result<S::Ok, S::Error>
104where
105    S: Serializer,
106{
107    let formatted = format!("{value:.N$}");
108    serializer.serialize_str(&formatted)
109}
110
111struct ReportData {
112    file_writers: RefCell<HashMap<TypeId, Writer<File>>>,
113    config: ConfigReportOptions,
114}
115
116// Registers a data container that stores
117// * file_writers: Maps report type to file writer
118// * config: Contains all the customizable filename options that the user supplies
119define_data_plugin!(
120    ReportPlugin,
121    ReportData,
122    ReportData {
123        file_writers: RefCell::new(HashMap::new()),
124        config: ConfigReportOptions::new(),
125    }
126);
127
128pub trait ContextReportExt: PluginContext {
129    // Builds the filename. Called by `add_report`, `short_name` refers to the
130    // report type. The three main components are `prefix`, `directory`, and
131    // `short_name`.
132    fn generate_filename(&mut self, short_name: &str) -> PathBuf {
133        let data_container = self.get_data_mut(ReportPlugin);
134        let prefix = &data_container.config.file_prefix;
135        let directory = &data_container.config.output_dir;
136        let short_name = short_name.to_string();
137        let basename = format!("{prefix}{short_name}");
138        directory.join(basename).with_extension("csv")
139    }
140
141    /// Add a report file keyed by a `TypeId`.
142    /// The `short_name` is used for file naming to distinguish what data each
143    /// output file points to.
144    /// # Errors
145    /// If the file already exists and `overwrite` is set to false, raises an error and info message.
146    /// If the file cannot be created, raises an error.
147    fn add_report_by_type_id(&mut self, type_id: TypeId, short_name: &str) -> Result<(), IxaError> {
148        trace!("adding report {short_name} by type_id {type_id:?}");
149        let path = self.generate_filename(short_name);
150
151        let data_container = self.get_data_mut(ReportPlugin);
152
153        let file_creation_result = File::create_new(&path);
154        let created_file = match file_creation_result {
155            Ok(file) => file,
156            Err(e) => match e.kind() {
157                std::io::ErrorKind::AlreadyExists => {
158                    if data_container.config.overwrite {
159                        File::create(&path)?
160                    } else {
161                        error!("File already exists: {}. Please set `overwrite` to true in the file configuration and rerun.", path.display());
162                        return Err(IxaError::IoError(e));
163                    }
164                }
165                _ => {
166                    return Err(IxaError::IoError(e));
167                }
168            },
169        };
170        let writer = Writer::from_writer(created_file);
171        let mut file_writer = data_container.file_writers.borrow_mut();
172        file_writer.insert(type_id, writer);
173        Ok(())
174    }
175
176    /// Call `add_report` with each report type, passing the name of the report type.
177    /// The `short_name` is used for file naming to distinguish what data each
178    /// output file points to.
179    /// # Errors
180    /// If the file already exists and `overwrite` is set to false, raises an error and info message.
181    /// If the file cannot be created, raises an error.
182    fn add_report<T: Report + 'static>(&mut self, short_name: &str) -> Result<(), IxaError> {
183        trace!("Adding report {short_name}");
184        self.add_report_by_type_id(TypeId::of::<T>(), short_name)
185    }
186
187    /// Adds a periodic report at the end of period `period` which summarizes the
188    /// number of people in each combination of properties in `tabulator`.
189    /// # Errors
190    /// If the file already exists and `overwrite` is set to false, raises an error and info message.
191    /// If the file cannot be created, returns [`IxaError`]
192    fn add_periodic_report<T: Tabulator + Clone + 'static>(
193        &mut self,
194        short_name: &str,
195        period: f64,
196        tabulator: T,
197    ) -> Result<(), IxaError> {
198        trace!("Adding periodic report {short_name}");
199
200        self.add_report_by_type_id(TypeId::of::<T>(), short_name)?;
201
202        {
203            // Write the header
204            let mut writer = self.get_writer(TypeId::of::<T>());
205            let columns = tabulator.get_columns();
206            let mut header = vec!["t".to_string()];
207            header.extend(columns);
208            header.push("count".to_string());
209            writer
210                .write_record(&header)
211                .expect("Failed to write header");
212        }
213
214        self.add_periodic_plan_with_phase(
215            period,
216            move |context: &mut Context| {
217                context.tabulate_person_properties(&tabulator, move |context, values, count| {
218                    let mut writer = context.get_writer(TypeId::of::<T>());
219                    let mut row = vec![context.get_current_time().to_string()];
220                    row.extend(values.to_owned());
221                    row.push(count.to_string());
222
223                    writer.write_record(&row).expect("Failed to write row");
224                });
225            },
226            crate::context::ExecutionPhase::Last,
227        );
228
229        Ok(())
230    }
231
232    fn get_writer(&self, type_id: TypeId) -> RefMut<Writer<File>> {
233        // No data container will exist if no reports have been added
234        let data_container = self.get_data(ReportPlugin);
235        let writers = data_container.file_writers.try_borrow_mut().unwrap();
236        RefMut::map(writers, |writers| {
237            writers
238                .get_mut(&type_id)
239                .expect("No writer found for the report type")
240        })
241    }
242
243    /// Write a new row to the appropriate report file
244    fn send_report<T: Report>(&self, report: T) {
245        let writer = &mut self.get_writer(report.type_id());
246        report.serialize(writer);
247    }
248
249    /// Returns a `ConfigReportOptions` object which has setter methods for report configuration
250    fn report_options(&mut self) -> &mut ConfigReportOptions {
251        let data_container = self.get_data_mut(ReportPlugin);
252        &mut data_container.config
253    }
254}
255impl ContextReportExt for Context {}
256
257#[cfg(test)]
258mod test {
259    use crate::{define_person_property_with_default, info};
260
261    use super::*;
262    use core::convert::TryInto;
263    use serde_derive::{Deserialize, Serialize};
264    use std::thread;
265    use tempfile::tempdir;
266
267    define_person_property_with_default!(IsRunner, bool, false);
268
269    #[derive(Serialize, Deserialize)]
270    struct SampleReport {
271        id: u32,
272        value: String,
273    }
274
275    define_report!(SampleReport);
276
277    #[test]
278    fn add_and_send_report() {
279        let temp_dir = tempdir().unwrap();
280        let path = PathBuf::from(&temp_dir.path());
281        // We need the writer to go out of scope so the file is flushed
282        {
283            let mut context = Context::new();
284            let config = context.report_options();
285            config
286                .file_prefix("prefix1_".to_string())
287                .directory(path.clone());
288            context.add_report::<SampleReport>("sample_report").unwrap();
289            let report = SampleReport {
290                id: 1,
291                value: "Test Value".to_string(),
292            };
293
294            context.send_report(report);
295        }
296
297        let file_path = path.join("prefix1_sample_report.csv");
298        assert!(file_path.exists(), "CSV file should exist");
299        assert!(file_path.metadata().unwrap().len() > 0);
300
301        let mut reader = csv::Reader::from_path(file_path).unwrap();
302        for result in reader.deserialize() {
303            let record: SampleReport = result.unwrap();
304            assert_eq!(record.id, 1);
305            assert_eq!(record.value, "Test Value");
306        }
307    }
308
309    #[test]
310    fn add_report_empty_prefix() {
311        let temp_dir = tempdir().unwrap();
312        let path = PathBuf::from(&temp_dir.path());
313        // We need the writer to go out of scope so the file is flushed
314        {
315            let mut context = Context::new();
316            let config = context.report_options();
317            config.directory(path.clone());
318            context.add_report::<SampleReport>("sample_report").unwrap();
319            let report = SampleReport {
320                id: 1,
321                value: "Test Value".to_string(),
322            };
323
324            context.send_report(report);
325        }
326        let file_path = path.join("sample_report.csv");
327        assert!(file_path.exists(), "CSV file should exist");
328        assert!(file_path.metadata().unwrap().len() > 0);
329
330        let mut reader = csv::Reader::from_path(file_path).unwrap();
331        for result in reader.deserialize() {
332            let record: SampleReport = result.unwrap();
333            assert_eq!(record.id, 1);
334            assert_eq!(record.value, "Test Value");
335        }
336    }
337
338    struct PathBufWithDrop {
339        file: PathBuf,
340    }
341
342    impl Drop for PathBufWithDrop {
343        fn drop(&mut self) {
344            std::fs::remove_file(&self.file).unwrap();
345        }
346    }
347
348    #[test]
349    fn add_report_no_dir() {
350        // We need the writer to go out of scope so the file is flushed
351        {
352            let mut context = Context::new();
353            let config = context.report_options();
354            config.file_prefix("test_prefix_".to_string());
355            context.add_report::<SampleReport>("sample_report").unwrap();
356            let report = SampleReport {
357                id: 1,
358                value: "Test Value".to_string(),
359            };
360
361            context.send_report(report);
362        }
363
364        let path = env::current_dir().unwrap();
365        let file_path = PathBufWithDrop {
366            file: path.join("test_prefix_sample_report.csv"),
367        };
368        assert!(file_path.file.exists(), "CSV file should exist");
369        assert!(file_path.file.metadata().unwrap().len() > 0);
370
371        let mut reader = csv::Reader::from_path(&file_path.file).unwrap();
372        for result in reader.deserialize() {
373            let record: SampleReport = result.unwrap();
374            assert_eq!(record.id, 1);
375            assert_eq!(record.value, "Test Value");
376        }
377    }
378
379    #[test]
380    #[should_panic(expected = "No writer found for the report type")]
381    fn send_report_without_adding_report() {
382        let context = Context::new();
383        let report = SampleReport {
384            id: 1,
385            value: "Test Value".to_string(),
386        };
387
388        context.send_report(report);
389    }
390
391    #[test]
392    fn multiple_reports_one_context() {
393        let temp_dir = tempdir().unwrap();
394        let path = PathBuf::from(&temp_dir.path());
395        // We need the writer to go out of scope so the file is flushed
396        {
397            let mut context = Context::new();
398            let config = context.report_options();
399            config
400                .file_prefix("mult_report_".to_string())
401                .directory(path.clone());
402            context.add_report::<SampleReport>("sample_report").unwrap();
403            let report1 = SampleReport {
404                id: 1,
405                value: "Value,1".to_string(),
406            };
407            let report2 = SampleReport {
408                id: 2,
409                value: "Value\n2".to_string(),
410            };
411
412            context.send_report(report1);
413            context.send_report(report2);
414        }
415
416        let file_path = path.join("mult_report_sample_report.csv");
417        assert!(file_path.exists(), "CSV file should exist");
418
419        let mut reader = csv::Reader::from_path(file_path).expect("Failed to open CSV file");
420        let mut records = reader.deserialize::<SampleReport>();
421
422        let item1: SampleReport = records
423            .next()
424            .expect("No record found")
425            .expect("Failed to deserialize record");
426        assert_eq!(item1.id, 1);
427        assert_eq!(item1.value, "Value,1");
428
429        let item2: SampleReport = records
430            .next()
431            .expect("No second record found")
432            .expect("Failed to deserialize record");
433        assert_eq!(item2.id, 2);
434        assert_eq!(item2.value, "Value\n2");
435    }
436
437    #[test]
438    fn multithreaded_report_generation_thread_local() {
439        let num_threads = 10;
440        let num_reports_per_thread = 5;
441
442        let mut handles = vec![];
443        let temp_dir = tempdir().unwrap();
444        let base_path = temp_dir.path().to_path_buf();
445
446        for i in 0..num_threads {
447            let path = base_path.clone();
448            let handle = thread::spawn(move || {
449                let mut context = Context::new();
450                let config = context.report_options();
451                config.file_prefix(i.to_string()).directory(path);
452                context.add_report::<SampleReport>("sample_report").unwrap();
453
454                for j in 0..num_reports_per_thread {
455                    let report = SampleReport {
456                        id: u32::try_from(i * num_reports_per_thread + j).unwrap(),
457                        value: format!("Thread {i} Report {j}"),
458                    };
459                    context.send_report(report);
460                }
461            });
462
463            handles.push(handle);
464        }
465
466        for handle in handles {
467            handle.join().expect("Thread failed");
468        }
469
470        for i in 0..num_threads {
471            let file_name = format!("{i}sample_report.csv");
472            let file_path = base_path.join(file_name);
473            assert!(file_path.exists(), "CSV file should exist");
474
475            let mut reader = csv::Reader::from_path(file_path).expect("Failed to open CSV file");
476            let records = reader.deserialize::<SampleReport>();
477
478            for (j, record) in records.enumerate() {
479                let record: SampleReport = record.expect("Failed to deserialize record");
480                let id_expected = TryInto::<u32>::try_into(i * num_reports_per_thread + j).unwrap();
481                assert_eq!(record.id, id_expected);
482            }
483        }
484    }
485
486    #[test]
487    fn dont_overwrite_report() {
488        let mut context1 = Context::new();
489        let temp_dir = tempdir().unwrap();
490        let path = PathBuf::from(&temp_dir.path());
491        let config = context1.report_options();
492        config
493            .file_prefix("prefix1_".to_string())
494            .directory(path.clone());
495        context1
496            .add_report::<SampleReport>("sample_report")
497            .unwrap();
498        let report = SampleReport {
499            id: 1,
500            value: "Test Value".to_string(),
501        };
502
503        context1.send_report(report);
504
505        let file_path = path.join("prefix1_sample_report.csv");
506        assert!(file_path.exists(), "CSV file should exist");
507
508        let mut context2 = Context::new();
509        let config = context2.report_options();
510        config.file_prefix("prefix1_".to_string()).directory(path);
511        info!("The next 'file already exists' error is intended for a passing test.");
512        let result = context2.add_report::<SampleReport>("sample_report");
513        assert!(result.is_err());
514        let error = result.err().unwrap();
515        match error {
516            IxaError::IoError(e) => {
517                assert_eq!(e.kind(), std::io::ErrorKind::AlreadyExists);
518            }
519            _ => {
520                panic!("Unexpected error type");
521            }
522        }
523    }
524
525    #[test]
526    fn overwrite_report() {
527        let mut context1 = Context::new();
528        let temp_dir = tempdir().unwrap();
529        let path = PathBuf::from(&temp_dir.path());
530        let config = context1.report_options();
531        config
532            .file_prefix("prefix1_".to_string())
533            .directory(path.clone());
534        context1
535            .add_report::<SampleReport>("sample_report")
536            .unwrap();
537        let report = SampleReport {
538            id: 1,
539            value: "Test Value".to_string(),
540        };
541
542        context1.send_report(report);
543
544        let file_path = path.join("prefix1_sample_report.csv");
545        assert!(file_path.exists(), "CSV file should exist");
546
547        let mut context2 = Context::new();
548        let config = context2.report_options();
549        config
550            .file_prefix("prefix1_".to_string())
551            .directory(path)
552            .overwrite(true);
553        let result = context2.add_report::<SampleReport>("sample_report");
554        assert!(result.is_ok());
555        let file = File::open(file_path).unwrap();
556        let reader = csv::Reader::from_reader(file);
557        let records = reader.into_records();
558        assert_eq!(records.count(), 0);
559    }
560
561    #[derive(PartialEq, Copy, Clone, Debug, Serialize, Deserialize)]
562    pub enum SymptomValue {
563        Presymptomatic,
564        Category1,
565        Category2,
566        Category3,
567        Category4,
568    }
569
570    define_person_property_with_default!(Symptoms, Option<SymptomValue>, None);
571
572    #[test]
573    fn add_periodic_report() {
574        let temp_dir = tempdir().unwrap();
575        let path = PathBuf::from(&temp_dir.path());
576        // We need the writer to go out of scope so the file is flushed
577        {
578            let mut context = Context::new();
579            let config = context.report_options();
580            config
581                .file_prefix("test_".to_string())
582                .directory(path.clone());
583            let _ = context.add_periodic_report("periodic", 1.2, (IsRunner, Symptoms));
584            let person = context.add_person(()).unwrap();
585            context.add_person(()).unwrap();
586
587            context.add_plan(1.2, move |context: &mut Context| {
588                context.set_person_property(person, IsRunner, true);
589                context.set_person_property(person, Symptoms, Some(SymptomValue::Category1));
590            });
591            context.execute();
592        }
593        let file_path = path.join("test_periodic.csv");
594        assert!(file_path.exists(), "CSV file should exist");
595
596        let mut reader = csv::Reader::from_path(file_path).unwrap();
597
598        assert_eq!(
599            reader.headers().unwrap(),
600            vec!["t", "IsRunner", "Symptoms", "count"]
601        );
602
603        let mut actual: Vec<Vec<String>> = reader
604            .records()
605            .map(|result| result.unwrap().iter().map(String::from).collect())
606            .collect();
607        let mut expected = vec![
608            vec!["0", "false", "None", "2"],
609            vec!["1.2", "false", "Category1", "0"],
610            vec!["1.2", "false", "None", "1"],
611            vec!["1.2", "true", "Category1", "1"],
612            vec!["1.2", "true", "None", "0"],
613        ];
614
615        actual.sort();
616        expected.sort();
617
618        assert_eq!(actual, expected, "CSV file should contain the correct data");
619    }
620}
ixa/report.rs

ixa/
report.rs
