ixa/
report.rs

1use crate::context::Context;
2use crate::error::IxaError;
3use crate::people::ContextPeopleExt;
4use crate::Tabulator;
5use crate::{error, trace};
6use crate::{HashMap, HashMapExt};
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! create_report_trait {
74    ($name:ident) => {
75        impl 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
119crate::context::define_data_plugin!(
120    ReportPlugin,
121    ReportData,
122    ReportData {
123        file_writers: RefCell::new(HashMap::new()),
124        config: ConfigReportOptions::new(),
125    }
126);
127
128impl Context {
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_container_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
142pub trait ContextReportExt {
143    /// Add a report file keyed by a `TypeId`.
144    /// The `short_name` is used for file naming to distinguish what data each
145    /// output file points to.
146    /// # Errors
147    /// If the file already exists and `overwrite` is set to false, raises an error and info message.
148    /// If the file cannot be created, raises an error.
149    fn add_report_by_type_id(&mut self, type_id: TypeId, short_name: &str) -> Result<(), IxaError>;
150
151    /// Call `add_report` with each report type, passing the name of the report type.
152    /// The `short_name` is used for file naming to distinguish what data each
153    /// output file points to.
154    /// # Errors
155    /// If the file already exists and `overwrite` is set to false, raises an error and info message.
156    /// If the file cannot be created, raises an error.
157    fn add_report<T: Report + 'static>(&mut self, short_name: &str) -> Result<(), IxaError>;
158
159    /// Adds a periodic report at the end of period `period` which summarizes the
160    /// number of people in each combination of properties in `tabulator`.
161    /// # Errors
162    /// If the file already exists and `overwrite` is set to false, raises an error and info message.
163    /// If the file cannot be created, returns [`IxaError`]
164    fn add_periodic_report<T: Tabulator + Clone + 'static>(
165        &mut self,
166        short_name: &str,
167        period: f64,
168        tabulator: T,
169    ) -> Result<(), IxaError>;
170    fn get_writer(&self, type_id: TypeId) -> RefMut<Writer<File>>;
171    fn send_report<T: Report>(&self, report: T);
172    fn report_options(&mut self) -> &mut ConfigReportOptions;
173}
174
175impl ContextReportExt for Context {
176    fn add_report_by_type_id(&mut self, type_id: TypeId, short_name: &str) -> Result<(), IxaError> {
177        trace!("adding report {short_name} by type_id {type_id:?}");
178        let path = self.generate_filename(short_name);
179
180        let data_container = self.get_data_container_mut(ReportPlugin);
181
182        let file_creation_result = File::create_new(&path);
183        let created_file = match file_creation_result {
184            Ok(file) => file,
185            Err(e) => match e.kind() {
186                std::io::ErrorKind::AlreadyExists => {
187                    if data_container.config.overwrite {
188                        File::create(&path)?
189                    } else {
190                        error!("File already exists: {}. Please set `overwrite` to true in the file configuration and rerun.", path.display());
191                        return Err(IxaError::IoError(e));
192                    }
193                }
194                _ => {
195                    return Err(IxaError::IoError(e));
196                }
197            },
198        };
199        let writer = Writer::from_writer(created_file);
200        let mut file_writer = data_container.file_writers.borrow_mut();
201        file_writer.insert(type_id, writer);
202        Ok(())
203    }
204    fn add_report<T: Report + 'static>(&mut self, short_name: &str) -> Result<(), IxaError> {
205        trace!("Adding report {short_name}");
206        self.add_report_by_type_id(TypeId::of::<T>(), short_name)
207    }
208    fn add_periodic_report<T: Tabulator + Clone + 'static>(
209        &mut self,
210        short_name: &str,
211        period: f64,
212        tabulator: T,
213    ) -> Result<(), IxaError> {
214        trace!("Adding periodic report {short_name}");
215
216        self.add_report_by_type_id(TypeId::of::<T>(), short_name)?;
217
218        {
219            // Write the header
220            let mut writer = self.get_writer(TypeId::of::<T>());
221            let columns = tabulator.get_columns();
222            let mut header = vec!["t".to_string()];
223            header.extend(columns);
224            header.push("count".to_string());
225            writer
226                .write_record(&header)
227                .expect("Failed to write header");
228        }
229
230        self.add_periodic_plan_with_phase(
231            period,
232            move |context: &mut Context| {
233                context.tabulate_person_properties(&tabulator, move |context, values, count| {
234                    let mut writer = context.get_writer(TypeId::of::<T>());
235                    let mut row = vec![context.get_current_time().to_string()];
236                    row.extend(values.to_owned());
237                    row.push(count.to_string());
238
239                    writer.write_record(&row).expect("Failed to write row");
240                });
241            },
242            crate::context::ExecutionPhase::Last,
243        );
244
245        Ok(())
246    }
247
248    fn get_writer(&self, type_id: TypeId) -> RefMut<Writer<File>> {
249        // No data container will exist if no reports have been added
250        let data_container = self
251            .get_data_container(ReportPlugin)
252            .expect("No writer found for the report type");
253        let writers = data_container.file_writers.try_borrow_mut().unwrap();
254        RefMut::map(writers, |writers| {
255            writers
256                .get_mut(&type_id)
257                .expect("No writer found for the report type")
258        })
259    }
260
261    /// Write a new row to the appropriate report file
262    fn send_report<T: Report>(&self, report: T) {
263        let writer = &mut self.get_writer(report.type_id());
264        report.serialize(writer);
265    }
266
267    /// Returns a `ConfigReportOptions` object which has setter methods for report configuration
268    fn report_options(&mut self) -> &mut ConfigReportOptions {
269        let data_container = self.get_data_container_mut(ReportPlugin);
270        &mut data_container.config
271    }
272}
273
274#[cfg(test)]
275mod test {
276    use crate::{define_person_property_with_default, info};
277
278    use super::*;
279    use core::convert::TryInto;
280    use serde_derive::{Deserialize, Serialize};
281    use std::thread;
282    use tempfile::tempdir;
283
284    define_person_property_with_default!(IsRunner, bool, false);
285
286    #[derive(Serialize, Deserialize)]
287    struct SampleReport {
288        id: u32,
289        value: String,
290    }
291
292    create_report_trait!(SampleReport);
293
294    #[test]
295    fn add_and_send_report() {
296        let temp_dir = tempdir().unwrap();
297        let path = PathBuf::from(&temp_dir.path());
298        // We need the writer to go out of scope so the file is flushed
299        {
300            let mut context = Context::new();
301            let config = context.report_options();
302            config
303                .file_prefix("prefix1_".to_string())
304                .directory(path.clone());
305            context.add_report::<SampleReport>("sample_report").unwrap();
306            let report = SampleReport {
307                id: 1,
308                value: "Test Value".to_string(),
309            };
310
311            context.send_report(report);
312        }
313
314        let file_path = path.join("prefix1_sample_report.csv");
315        assert!(file_path.exists(), "CSV file should exist");
316        assert!(file_path.metadata().unwrap().len() > 0);
317
318        let mut reader = csv::Reader::from_path(file_path).unwrap();
319        for result in reader.deserialize() {
320            let record: SampleReport = result.unwrap();
321            assert_eq!(record.id, 1);
322            assert_eq!(record.value, "Test Value");
323        }
324    }
325
326    #[test]
327    fn add_report_empty_prefix() {
328        let temp_dir = tempdir().unwrap();
329        let path = PathBuf::from(&temp_dir.path());
330        // We need the writer to go out of scope so the file is flushed
331        {
332            let mut context = Context::new();
333            let config = context.report_options();
334            config.directory(path.clone());
335            context.add_report::<SampleReport>("sample_report").unwrap();
336            let report = SampleReport {
337                id: 1,
338                value: "Test Value".to_string(),
339            };
340
341            context.send_report(report);
342        }
343        let file_path = path.join("sample_report.csv");
344        assert!(file_path.exists(), "CSV file should exist");
345        assert!(file_path.metadata().unwrap().len() > 0);
346
347        let mut reader = csv::Reader::from_path(file_path).unwrap();
348        for result in reader.deserialize() {
349            let record: SampleReport = result.unwrap();
350            assert_eq!(record.id, 1);
351            assert_eq!(record.value, "Test Value");
352        }
353    }
354
355    struct PathBufWithDrop {
356        file: PathBuf,
357    }
358
359    impl Drop for PathBufWithDrop {
360        fn drop(&mut self) {
361            std::fs::remove_file(&self.file).unwrap();
362        }
363    }
364
365    #[test]
366    fn add_report_no_dir() {
367        // We need the writer to go out of scope so the file is flushed
368        {
369            let mut context = Context::new();
370            let config = context.report_options();
371            config.file_prefix("test_prefix_".to_string());
372            context.add_report::<SampleReport>("sample_report").unwrap();
373            let report = SampleReport {
374                id: 1,
375                value: "Test Value".to_string(),
376            };
377
378            context.send_report(report);
379        }
380
381        let path = env::current_dir().unwrap();
382        let file_path = PathBufWithDrop {
383            file: path.join("test_prefix_sample_report.csv"),
384        };
385        assert!(file_path.file.exists(), "CSV file should exist");
386        assert!(file_path.file.metadata().unwrap().len() > 0);
387
388        let mut reader = csv::Reader::from_path(&file_path.file).unwrap();
389        for result in reader.deserialize() {
390            let record: SampleReport = result.unwrap();
391            assert_eq!(record.id, 1);
392            assert_eq!(record.value, "Test Value");
393        }
394    }
395
396    #[test]
397    #[should_panic(expected = "No writer found for the report type")]
398    fn send_report_without_adding_report() {
399        let context = Context::new();
400        let report = SampleReport {
401            id: 1,
402            value: "Test Value".to_string(),
403        };
404
405        context.send_report(report);
406    }
407
408    #[test]
409    fn multiple_reports_one_context() {
410        let temp_dir = tempdir().unwrap();
411        let path = PathBuf::from(&temp_dir.path());
412        // We need the writer to go out of scope so the file is flushed
413        {
414            let mut context = Context::new();
415            let config = context.report_options();
416            config
417                .file_prefix("mult_report_".to_string())
418                .directory(path.clone());
419            context.add_report::<SampleReport>("sample_report").unwrap();
420            let report1 = SampleReport {
421                id: 1,
422                value: "Value,1".to_string(),
423            };
424            let report2 = SampleReport {
425                id: 2,
426                value: "Value\n2".to_string(),
427            };
428
429            context.send_report(report1);
430            context.send_report(report2);
431        }
432
433        let file_path = path.join("mult_report_sample_report.csv");
434        assert!(file_path.exists(), "CSV file should exist");
435
436        let mut reader = csv::Reader::from_path(file_path).expect("Failed to open CSV file");
437        let mut records = reader.deserialize::<SampleReport>();
438
439        let item1: SampleReport = records
440            .next()
441            .expect("No record found")
442            .expect("Failed to deserialize record");
443        assert_eq!(item1.id, 1);
444        assert_eq!(item1.value, "Value,1");
445
446        let item2: SampleReport = records
447            .next()
448            .expect("No second record found")
449            .expect("Failed to deserialize record");
450        assert_eq!(item2.id, 2);
451        assert_eq!(item2.value, "Value\n2");
452    }
453
454    #[test]
455    fn multithreaded_report_generation_thread_local() {
456        let num_threads = 10;
457        let num_reports_per_thread = 5;
458
459        let mut handles = vec![];
460        let temp_dir = tempdir().unwrap();
461        let base_path = temp_dir.path().to_path_buf();
462
463        for i in 0..num_threads {
464            let path = base_path.clone();
465            let handle = thread::spawn(move || {
466                let mut context = Context::new();
467                let config = context.report_options();
468                config.file_prefix(i.to_string()).directory(path);
469                context.add_report::<SampleReport>("sample_report").unwrap();
470
471                for j in 0..num_reports_per_thread {
472                    let report = SampleReport {
473                        id: u32::try_from(i * num_reports_per_thread + j).unwrap(),
474                        value: format!("Thread {i} Report {j}"),
475                    };
476                    context.send_report(report);
477                }
478            });
479
480            handles.push(handle);
481        }
482
483        for handle in handles {
484            handle.join().expect("Thread failed");
485        }
486
487        for i in 0..num_threads {
488            let file_name = format!("{i}sample_report.csv");
489            let file_path = base_path.join(file_name);
490            assert!(file_path.exists(), "CSV file should exist");
491
492            let mut reader = csv::Reader::from_path(file_path).expect("Failed to open CSV file");
493            let records = reader.deserialize::<SampleReport>();
494
495            for (j, record) in records.enumerate() {
496                let record: SampleReport = record.expect("Failed to deserialize record");
497                let id_expected = TryInto::<u32>::try_into(i * num_reports_per_thread + j).unwrap();
498                assert_eq!(record.id, id_expected);
499            }
500        }
501    }
502
503    #[test]
504    fn dont_overwrite_report() {
505        let mut context1 = Context::new();
506        let temp_dir = tempdir().unwrap();
507        let path = PathBuf::from(&temp_dir.path());
508        let config = context1.report_options();
509        config
510            .file_prefix("prefix1_".to_string())
511            .directory(path.clone());
512        context1
513            .add_report::<SampleReport>("sample_report")
514            .unwrap();
515        let report = SampleReport {
516            id: 1,
517            value: "Test Value".to_string(),
518        };
519
520        context1.send_report(report);
521
522        let file_path = path.join("prefix1_sample_report.csv");
523        assert!(file_path.exists(), "CSV file should exist");
524
525        let mut context2 = Context::new();
526        let config = context2.report_options();
527        config.file_prefix("prefix1_".to_string()).directory(path);
528        info!("The next 'file already exists' error is intended for a passing test.");
529        let result = context2.add_report::<SampleReport>("sample_report");
530        assert!(result.is_err());
531        let error = result.err().unwrap();
532        match error {
533            IxaError::IoError(e) => {
534                assert_eq!(e.kind(), std::io::ErrorKind::AlreadyExists);
535            }
536            _ => {
537                panic!("Unexpected error type");
538            }
539        }
540    }
541
542    #[test]
543    fn overwrite_report() {
544        let mut context1 = Context::new();
545        let temp_dir = tempdir().unwrap();
546        let path = PathBuf::from(&temp_dir.path());
547        let config = context1.report_options();
548        config
549            .file_prefix("prefix1_".to_string())
550            .directory(path.clone());
551        context1
552            .add_report::<SampleReport>("sample_report")
553            .unwrap();
554        let report = SampleReport {
555            id: 1,
556            value: "Test Value".to_string(),
557        };
558
559        context1.send_report(report);
560
561        let file_path = path.join("prefix1_sample_report.csv");
562        assert!(file_path.exists(), "CSV file should exist");
563
564        let mut context2 = Context::new();
565        let config = context2.report_options();
566        config
567            .file_prefix("prefix1_".to_string())
568            .directory(path)
569            .overwrite(true);
570        let result = context2.add_report::<SampleReport>("sample_report");
571        assert!(result.is_ok());
572        let file = File::open(file_path).unwrap();
573        let reader = csv::Reader::from_reader(file);
574        let records = reader.into_records();
575        assert_eq!(records.count(), 0);
576    }
577
578    #[derive(PartialEq, Copy, Clone, Debug, Serialize, Deserialize)]
579    pub enum SymptomValue {
580        Presymptomatic,
581        Category1,
582        Category2,
583        Category3,
584        Category4,
585    }
586
587    define_person_property_with_default!(Symptoms, Option<SymptomValue>, None);
588
589    #[test]
590    fn add_periodic_report() {
591        let temp_dir = tempdir().unwrap();
592        let path = PathBuf::from(&temp_dir.path());
593        // We need the writer to go out of scope so the file is flushed
594        {
595            let mut context = Context::new();
596            let config = context.report_options();
597            config
598                .file_prefix("test_".to_string())
599                .directory(path.clone());
600            let _ = context.add_periodic_report("periodic", 1.2, (IsRunner, Symptoms));
601            let person = context.add_person(()).unwrap();
602            context.add_person(()).unwrap();
603
604            context.add_plan(1.2, move |context: &mut Context| {
605                context.set_person_property(person, IsRunner, true);
606                context.set_person_property(person, Symptoms, Some(SymptomValue::Category1));
607            });
608            context.execute();
609        }
610        let file_path = path.join("test_periodic.csv");
611        assert!(file_path.exists(), "CSV file should exist");
612
613        let mut reader = csv::Reader::from_path(file_path).unwrap();
614
615        assert_eq!(
616            reader.headers().unwrap(),
617            vec!["t", "IsRunner", "Symptoms", "count"]
618        );
619
620        let mut actual: Vec<Vec<String>> = reader
621            .records()
622            .map(|result| result.unwrap().iter().map(String::from).collect())
623            .collect();
624        let mut expected = vec![
625            vec!["0", "false", "None", "2"],
626            vec!["1.2", "false", "Category1", "0"],
627            vec!["1.2", "false", "None", "1"],
628            vec!["1.2", "true", "Category1", "1"],
629            vec!["1.2", "true", "None", "0"],
630        ];
631
632        actual.sort();
633        expected.sort();
634
635        assert_eq!(actual, expected, "CSV file should contain the correct data");
636    }
637}
ixa/report.rs

ixa/
report.rs
