brk_query/
lib.rs

1#![doc = include_str!("../README.md")]
2#![doc = "\n## Example\n\n```rust"]
3#![doc = include_str!("../examples/main.rs")]
4#![doc = "```"]
5
6use brk_computer::Computer;
7use brk_indexer::Indexer;
8use brk_vec::AnyStorableVec;
9use tabled::settings::Style;
10
11mod format;
12mod index;
13mod output;
14mod params;
15mod table;
16mod vec_trees;
17
18pub use format::Format;
19pub use index::Index;
20pub use output::{Output, Value};
21pub use params::Params;
22pub use table::Tabled;
23use vec_trees::VecTrees;
24
25pub struct Query<'a> {
26    pub vec_trees: VecTrees<'a>,
27    _indexer: &'a Indexer,
28    _computer: &'a Computer,
29}
30
31impl<'a> Query<'a> {
32    pub fn build(indexer: &'a Indexer, computer: &'a Computer) -> Self {
33        let mut vec_trees = VecTrees::default();
34
35        indexer
36            .vecs()
37            .as_any_vecs()
38            .into_iter()
39            .for_each(|vec| vec_trees.insert(vec));
40
41        computer
42            .vecs()
43            .as_any_vecs()
44            .into_iter()
45            .for_each(|vec| vec_trees.insert(vec));
46
47        Self {
48            vec_trees,
49            _indexer: indexer,
50            _computer: computer,
51        }
52    }
53
54    pub fn search(&self, index: Index, ids: &[&str]) -> Vec<(String, &&dyn AnyStorableVec)> {
55        let tuples = ids
56            .iter()
57            .flat_map(|s| {
58                s.to_lowercase()
59                    .replace("_", "-")
60                    .split_whitespace()
61                    .flat_map(|s| {
62                        s.split(',')
63                            .flat_map(|s| s.split('+').map(|s| s.to_string()))
64                    })
65                    .collect::<Vec<_>>()
66            })
67            .map(|mut id| {
68                let mut res = self.vec_trees.id_to_index_to_vec.get(&id);
69                if res.is_none() {
70                    if let Ok(index) = Index::try_from(id.as_str()) {
71                        id = index.possible_values().last().unwrap().to_string();
72                        res = self.vec_trees.id_to_index_to_vec.get(&id)
73                    }
74                }
75                (id, res)
76            })
77            .filter(|(_, opt)| opt.is_some())
78            .map(|(id, vec)| (id, vec.unwrap()))
79            .collect::<Vec<_>>();
80
81        tuples
82            .iter()
83            .flat_map(|(str, i_to_v)| i_to_v.get(&index).map(|vec| (str.to_owned(), vec)))
84            .collect::<Vec<_>>()
85    }
86
87    pub fn format(
88        &self,
89        vecs: Vec<(String, &&dyn AnyStorableVec)>,
90        from: Option<i64>,
91        to: Option<i64>,
92        format: Option<Format>,
93    ) -> color_eyre::Result<Output> {
94        let mut values = vecs
95            .iter()
96            .map(|(_, vec)| -> brk_vec::Result<Vec<serde_json::Value>> {
97                vec.collect_range_values(from, to)
98            })
99            .collect::<brk_vec::Result<Vec<_>>>()?;
100
101        if values.is_empty() {
102            return Ok(Output::default(format));
103        }
104
105        let ids_last_i = vecs.len() - 1;
106
107        Ok(match format {
108            Some(Format::CSV) | Some(Format::TSV) => {
109                let delimiter = if format == Some(Format::CSV) {
110                    ','
111                } else {
112                    '\t'
113                };
114
115                let mut text = vecs
116                    .iter()
117                    .map(|(id, _)| id.to_owned())
118                    .collect::<Vec<_>>()
119                    .join(&delimiter.to_string());
120
121                text.push('\n');
122
123                let values_len = values.first().unwrap().len();
124
125                (0..values_len).for_each(|i| {
126                    let mut line = "".to_string();
127                    values.iter().enumerate().for_each(|(id_i, v)| {
128                        line += &v.get(i).unwrap().to_string();
129                        if id_i == ids_last_i {
130                            line.push('\n');
131                        } else {
132                            line.push(delimiter);
133                        }
134                    });
135                    text += &line;
136                });
137
138                if format == Some(Format::CSV) {
139                    Output::CSV(text)
140                } else {
141                    Output::TSV(text)
142                }
143            }
144            Some(Format::MD) => {
145                let mut table =
146                    values.to_table(vecs.iter().map(|(s, _)| s.to_owned()).collect::<Vec<_>>());
147
148                table.with(Style::markdown());
149
150                Output::MD(table.to_string())
151            }
152            Some(Format::JSON) | None => {
153                if values.len() == 1 {
154                    let mut values = values.pop().unwrap();
155                    if values.len() == 1 {
156                        let value = values.pop().unwrap();
157                        Output::Json(Value::Single(value))
158                    } else {
159                        Output::Json(Value::List(values))
160                    }
161                } else {
162                    Output::Json(Value::Matrix(values))
163                }
164            }
165        })
166    }
167
168    pub fn search_and_format(
169        &self,
170        index: Index,
171        ids: &[&str],
172        from: Option<i64>,
173        to: Option<i64>,
174        format: Option<Format>,
175    ) -> color_eyre::Result<Output> {
176        self.format(self.search(index, ids), from, to, format)
177    }
178}
brk_query/lib.rs

brk_query/
lib.rs
