jawk/
extractor.rs

1use std::{io::Read, rc::Rc};
2
3use crate::{
4    json_value::JsonValue,
5    processor::Context,
6    reader::Reader,
7    selection::{Get, Result, SelectionParseError},
8};
9
10enum SingleExtract {
11    ByKey(String),
12    ByIndex(usize),
13}
14
15enum ExtractFromInput {
16    Root,
17    Element(Vec<SingleExtract>),
18}
19
20struct Extract {
21    number_of_parents: usize,
22    extract_from_input: ExtractFromInput,
23}
24
25impl ExtractFromInput {
26    fn extract(&self, input: &JsonValue) -> Option<JsonValue> {
27        match self {
28            ExtractFromInput::Root => Some(input.clone()),
29            ExtractFromInput::Element(es) => {
30                let mut val = Some(input.clone());
31                for e in es {
32                    match val {
33                        None => {
34                            break;
35                        }
36                        Some(value) => {
37                            val = e.extract(&value);
38                        }
39                    }
40                }
41                val
42            }
43        }
44    }
45}
46
47impl Get for Extract {
48    fn get(&self, value: &Context) -> Option<JsonValue> {
49        let input = value.parent_input(self.number_of_parents);
50        self.extract_from_input.extract(input)
51    }
52}
53
54pub fn parse_extractor<R: Read>(reader: &mut Reader<R>) -> Result<Rc<dyn Get>> {
55    let number_of_parents = read_number_of_parents(reader)?;
56    let extract_from_input = ExtractFromInput::parse(reader)?;
57    Ok(Rc::new(Extract {
58        number_of_parents,
59        extract_from_input,
60    }))
61}
62
63fn read_number_of_parents<R: Read>(reader: &mut Reader<R>) -> Result<usize> {
64    let mut size = 0;
65    loop {
66        match reader.peek()? {
67            Some(b'^') => {
68                size += 1;
69                reader.next()?;
70            }
71            _ => {
72                return Ok(size);
73            }
74        }
75    }
76}
77
78pub fn root() -> Rc<dyn Get> {
79    Rc::new(Extract {
80        extract_from_input: ExtractFromInput::Root,
81        number_of_parents: 0,
82    })
83}
84
85impl ExtractFromInput {
86    fn parse<R: Read>(reader: &mut Reader<R>) -> Result<Self> {
87        let mut ext = vec![];
88        loop {
89            match reader.peek()? {
90                Some(b'.') => {
91                    let key = Self::read_extract_key(reader)?;
92                    if key.is_empty() {
93                        return if ext.is_empty() {
94                            Ok(ExtractFromInput::Root)
95                        } else {
96                            Err(SelectionParseError::MissingKey(reader.where_am_i()))
97                        };
98                    }
99                    let es = SingleExtract::ByKey(key);
100                    ext.push(es);
101                }
102                Some(b'#') => match Self::read_extract_index(reader)? {
103                    None => {
104                        return if ext.is_empty() {
105                            Ok(ExtractFromInput::Root)
106                        } else {
107                            Err(SelectionParseError::MissingKey(reader.where_am_i()))
108                        };
109                    }
110                    Some(index) => {
111                        let es = SingleExtract::ByIndex(index);
112                        ext.push(es);
113                    }
114                },
115                _ => {
116                    return Ok(ExtractFromInput::Element(ext));
117                }
118            }
119        }
120    }
121    fn read_extract_key<R: Read>(reader: &mut Reader<R>) -> Result<String> {
122        let mut buf = Vec::new();
123        loop {
124            match reader.next()? {
125                None => {
126                    break;
127                }
128                Some(ch) => {
129                    if ch.is_ascii_whitespace()
130                        || ch == b'.'
131                        || ch == b','
132                        || ch == b'='
133                        || ch == b'('
134                        || ch == b')'
135                        || ch.is_ascii_control()
136                        || ch == b'\"'
137                        || ch == b']'
138                        || ch == b'['
139                        || ch == b'{'
140                        || ch == b'}'
141                        || ch == b'#'
142                    {
143                        break;
144                    }
145                    buf.push(ch);
146                }
147            }
148        }
149        let str = String::from_utf8(buf)?;
150        Ok(str)
151    }
152    fn read_extract_index<R: Read>(reader: &mut Reader<R>) -> Result<Option<usize>> {
153        reader.next()?;
154        let mut digits = Vec::new();
155        reader.read_digits(&mut digits)?;
156        let str = String::from_utf8(digits)?;
157        if str.is_empty() {
158            return Ok(None);
159        }
160        let number = str.parse::<usize>()?;
161        Ok(Some(number))
162    }
163}
164
165impl SingleExtract {
166    fn extract(&self, value: &JsonValue) -> Option<JsonValue> {
167        match value {
168            JsonValue::Array(list) => match self {
169                SingleExtract::ByIndex(index) => list.get(*index).cloned(),
170                _ => None,
171            },
172            JsonValue::Object(map) => match self {
173                SingleExtract::ByKey(key) => map.get(key).cloned(),
174                _ => None,
175            },
176            _ => None,
177        }
178    }
179}