1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

use crate::environment;
use std::collections::{HashMap, HashSet};

#[cfg(test)]
#[path = "./checkpoint_test.rs"]
mod checkpoint_test;

/// Environment variable checkpoint
///
/// A checkpoint stores all environment variables present at the time it was initialized,
/// which then can be restored using the [`restore`] method.
///
/// A checkpoint can be "impure", by excluding a set of variables, which should not be checkpointed,
/// if not specified or added this will reset **every** value.
pub struct Checkpoint {
    snapshot: HashMap<String, String>,
    exclude: HashSet<String>,
}

impl Checkpoint {
    /// Create a new Checkpoint
    ///
    /// This captures all currently present environment variables.
    pub fn new() -> Self {
        let mut env = HashMap::new();
        env.extend(environment::vars());

        Self {
            snapshot: env,
            exclude: HashSet::new(),
        }
    }

    /// Exclude a key from the [`Checkpoint::restore`]
    ///
    /// Keys that have been excluded will not be removed, created or modified in the restore process.
    pub fn exclude<T: Into<String>>(&mut self, key: T) -> &mut Self {
        self.exclude.insert(key.into());

        self
    }

    /// Restore the current state of environment to the state the checkpoint was taken in.
    ///
    /// This will remove added keys, re-create values and set new values.
    /// This will skip setting unchanged values.
    pub fn restore(self) {
        let chk = self.snapshot;

        let mut env = HashMap::new();
        env.extend(environment::vars());

        let exclude: HashSet<_> = self.exclude.iter().collect();

        // only change the variables that are of significance, this means we need to partition into
        // 3 piles:
        //  * to be removed
        //  * to be added
        //  * to be changed

        let chk_keys = &chk.keys().collect::<HashSet<_>>() - &exclude;
        let env_keys = &env.keys().collect::<HashSet<_>>() - &exclude;

        let remove = &env_keys - &chk_keys;
        let create = &chk_keys - &env_keys;

        let modify = (&chk_keys & &env_keys)
            .into_iter()
            .filter(|key| &chk[*key] != &env[*key])
            .collect::<HashSet<_>>();

        for key in remove {
            environment::remove(key);
        }

        for key in &create | &modify {
            environment::set(key, &chk[key]);
        }
    }
}