dotenvy/
iter.rs

Help
1use std::collections::HashMap;
2use std::env;
3use std::io::prelude::*;
4use std::io::BufReader;
5
6use crate::errors::*;
7use crate::parse;
8
9pub struct Iter<R> {
10    lines: QuotedLines<BufReader<R>>,
11    substitution_data: HashMap<String, Option<String>>,
12}
13
14impl<R: Read> Iter<R> {
15    pub fn new(reader: R) -> Iter<R> {
16        Iter {
17            lines: QuotedLines {
18                buf: BufReader::new(reader),
19            },
20            substitution_data: HashMap::new(),
21        }
22    }
23
24    /// Loads all variables found in the `reader` into the environment,
25    /// preserving any existing environment variables of the same name.
26    ///
27    /// If a variable is specified multiple times within the reader's data,
28    /// then the first occurrence is applied.
29    pub fn load(mut self) -> Result<()> {
30        self.remove_bom()?;
31
32        for item in self {
33            let (key, value) = item?;
34            if env::var(&key).is_err() {
35                env::set_var(&key, value);
36            }
37        }
38
39        Ok(())
40    }
41
42    /// Loads all variables found in the `reader` into the environment,
43    /// overriding any existing environment variables of the same name.
44    ///
45    /// If a variable is specified multiple times within the reader's data,
46    /// then the last occurrence is applied.
47    pub fn load_override(mut self) -> Result<()> {
48        self.remove_bom()?;
49
50        for item in self {
51            let (key, value) = item?;
52            env::set_var(key, value);
53        }
54
55        Ok(())
56    }
57
58    fn remove_bom(&mut self) -> Result<()> {
59        let buffer = self.lines.buf.fill_buf().map_err(Error::Io)?;
60        // https://www.compart.com/en/unicode/U+FEFF
61        if buffer.starts_with(&[0xEF, 0xBB, 0xBF]) {
62            // remove the BOM from the bufreader
63            self.lines.buf.consume(3);
64        }
65        Ok(())
66    }
67}
68
69struct QuotedLines<B> {
70    buf: B,
71}
72
73enum ParseState {
74    Complete,
75    Escape,
76    StrongOpen,
77    StrongOpenEscape,
78    WeakOpen,
79    WeakOpenEscape,
80    Comment,
81    WhiteSpace,
82}
83
84fn eval_end_state(prev_state: ParseState, buf: &str) -> (usize, ParseState) {
85    let mut cur_state = prev_state;
86    let mut cur_pos: usize = 0;
87
88    for (pos, c) in buf.char_indices() {
89        cur_pos = pos;
90        cur_state = match cur_state {
91            ParseState::WhiteSpace => match c {
92                '#' => return (cur_pos, ParseState::Comment),
93                '\\' => ParseState::Escape,
94                '"' => ParseState::WeakOpen,
95                '\'' => ParseState::StrongOpen,
96                _ => ParseState::Complete,
97            },
98            ParseState::Escape => ParseState::Complete,
99            ParseState::Complete => match c {
100                c if c.is_whitespace() && c != '\n' && c != '\r' => ParseState::WhiteSpace,
101                '\\' => ParseState::Escape,
102                '"' => ParseState::WeakOpen,
103                '\'' => ParseState::StrongOpen,
104                _ => ParseState::Complete,
105            },
106            ParseState::WeakOpen => match c {
107                '\\' => ParseState::WeakOpenEscape,
108                '"' => ParseState::Complete,
109                _ => ParseState::WeakOpen,
110            },
111            ParseState::WeakOpenEscape => ParseState::WeakOpen,
112            ParseState::StrongOpen => match c {
113                '\\' => ParseState::StrongOpenEscape,
114                '\'' => ParseState::Complete,
115                _ => ParseState::StrongOpen,
116            },
117            ParseState::StrongOpenEscape => ParseState::StrongOpen,
118            // Comments last the entire line.
119            ParseState::Comment => panic!("should have returned early"),
120        };
121    }
122    (cur_pos, cur_state)
123}
124
125impl<B: BufRead> Iterator for QuotedLines<B> {
126    type Item = Result<String>;
127
128    fn next(&mut self) -> Option<Result<String>> {
129        let mut buf = String::new();
130        let mut cur_state = ParseState::Complete;
131        let mut buf_pos;
132        let mut cur_pos;
133        loop {
134            buf_pos = buf.len();
135            match self.buf.read_line(&mut buf) {
136                Ok(0) => match cur_state {
137                    ParseState::Complete => return None,
138                    _ => {
139                        let len = buf.len();
140                        return Some(Err(Error::LineParse(buf, len)));
141                    }
142                },
143                Ok(_n) => {
144                    // Skip lines which start with a # before iteration
145                    // This optimizes parsing a bit.
146                    if buf.trim_start().starts_with('#') {
147                        return Some(Ok(String::with_capacity(0)));
148                    }
149                    let result = eval_end_state(cur_state, &buf[buf_pos..]);
150                    cur_pos = result.0;
151                    cur_state = result.1;
152
153                    match cur_state {
154                        ParseState::Complete => {
155                            if buf.ends_with('\n') {
156                                buf.pop();
157                                if buf.ends_with('\r') {
158                                    buf.pop();
159                                }
160                            }
161                            return Some(Ok(buf));
162                        }
163                        ParseState::Escape
164                        | ParseState::StrongOpen
165                        | ParseState::StrongOpenEscape
166                        | ParseState::WeakOpen
167                        | ParseState::WeakOpenEscape
168                        | ParseState::WhiteSpace => {}
169                        ParseState::Comment => {
170                            buf.truncate(buf_pos + cur_pos);
171                            return Some(Ok(buf));
172                        }
173                    }
174                }
175                Err(e) => return Some(Err(Error::Io(e))),
176            }
177        }
178    }
179}
180
181impl<R: Read> Iterator for Iter<R> {
182    type Item = Result<(String, String)>;
183
184    fn next(&mut self) -> Option<Self::Item> {
185        loop {
186            let line = match self.lines.next() {
187                Some(Ok(line)) => line,
188                Some(Err(err)) => return Some(Err(err)),
189                None => return None,
190            };
191
192            match parse::parse_line(&line, &mut self.substitution_data) {
193                Ok(Some(result)) => return Some(Ok(result)),
194                Ok(None) => {}
195                Err(err) => return Some(Err(err)),
196            }
197        }
198    }
199}
dotenvy/iter.rs

dotenvy/
iter.rs
