simd_csv/select/
dsl.rs

1/// This module contains the parsing logic of a tiny column selection DSL.
2///
3/// It comes directly from [`xan`](https://github.com/medialab/xan/) and was
4/// originally formulated and implemented by @BurntSushi for
5/// [`xsv`](https://github.com/burntsushi/xsv).
6///
7/// `xan` changed multiple things from the original `xsv` implementation:
8///     - indexation is now zero-based
9///     - range char was changed to `:` instead of `-`
10///     - added negative indexing
11///     - added wildcard selection
12use std::cmp::Ordering;
13use std::collections::BTreeMap;
14use std::convert::TryFrom;
15use std::fmt;
16use std::str::FromStr;
17
18use super::selection::Selection;
19use crate::error::{Error, ErrorKind};
20
21/// A parsed selector that can be applied on CSV headers to create a
22/// [`crate::Selection`].
23#[derive(Clone)]
24pub struct Selector {
25    selectors: Vec<CompositeSelector>,
26    invert: bool,
27}
28
29impl FromStr for Selector {
30    type Err = Error;
31
32    fn from_str(mut s: &str) -> Result<Self, Self::Err> {
33        let invert = if !s.is_empty() && s.as_bytes()[0] == b'!' {
34            s = &s[1..];
35            true
36        } else {
37            false
38        };
39        Ok(Self {
40            selectors: SelectorParser::new(s)
41                .parse()
42                .map_err(|msg| Error::new(ErrorKind::SelectorParseError(msg)))?,
43            invert,
44        })
45    }
46}
47
48impl Selector {
49    pub fn is_empty(&self) -> bool {
50        self.selectors.is_empty()
51    }
52
53    pub fn invert(&mut self) {
54        self.invert = !self.invert;
55    }
56
57    pub fn select<'a, H>(&self, first_record: H, use_names: bool) -> Result<Selection, Error>
58    where
59        H: IntoIterator<Item = &'a [u8]>,
60    {
61        let first_record = first_record.into_iter().collect::<Vec<_>>();
62
63        if self.selectors.is_empty() {
64            return Ok(Selection::new(
65                if self.invert {
66                    // Inverting everything means we get nothing.
67                    vec![]
68                } else {
69                    (0..first_record.len()).collect()
70                },
71                first_record.len(),
72            ));
73        }
74
75        let mut map = vec![];
76        for sel in &self.selectors {
77            let idxs = sel.indices(&first_record, use_names);
78            map.extend(idxs?.into_iter());
79        }
80        if self.invert {
81            let mut new_map = vec![];
82            for i in 0..first_record.len() {
83                if !map.contains(&i) {
84                    new_map.push(i);
85                }
86            }
87            return Ok(Selection::new(new_map, first_record.len()));
88        }
89        Ok(Selection::new(map, first_record.len()))
90    }
91
92    pub fn select_one<'a, H>(&self, first_record: H, use_names: bool) -> Result<usize, Error>
93    where
94        H: IntoIterator<Item = &'a [u8]>,
95    {
96        let selection = self.select(first_record, use_names)?;
97
98        if selection.len() != 1 {
99            return Err(Error::new(ErrorKind::SelectionError(
100                "target selection is not a single column".to_string(),
101            )));
102        }
103
104        Ok(selection[0])
105    }
106
107    pub fn retain_known<'a, H>(&mut self, headers: H) -> Vec<usize>
108    where
109        H: IntoIterator<Item = &'a [u8]>,
110    {
111        let headers = headers.into_iter().collect::<Vec<_>>();
112
113        let mut dropped: Vec<usize> = Vec::new();
114
115        for (i, selector) in self.selectors.iter().enumerate() {
116            match selector {
117                CompositeSelector::One(sel) if sel.index(&headers, true).is_err() => {
118                    dropped.push(i);
119                }
120                CompositeSelector::Range(start, end)
121                    if start.index(&headers, true).is_err()
122                        && end.index(&headers, true).is_err() =>
123                {
124                    dropped.push(i);
125                }
126                _ => continue,
127            };
128        }
129
130        let mut i: usize = 0;
131
132        self.selectors.retain(|_| {
133            let drop = !dropped.contains(&i);
134
135            i += 1;
136
137            drop
138        });
139
140        dropped
141    }
142}
143
144impl fmt::Debug for Selector {
145    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
146        if self.selectors.is_empty() {
147            write!(f, "<All>")
148        } else {
149            let strs: Vec<_> = self
150                .selectors
151                .iter()
152                .map(|sel| format!("{:?}", sel))
153                .collect();
154            write!(f, "{}", strs.join(", "))
155        }
156    }
157}
158
159impl TryFrom<String> for Selector {
160    type Error = Error;
161
162    fn try_from(value: String) -> Result<Self, Self::Error> {
163        value.parse()
164    }
165}
166
167impl Default for Selector {
168    fn default() -> Self {
169        "".parse().unwrap()
170    }
171}
172
173struct SelectorParser {
174    chars: Vec<char>,
175    pos: usize,
176}
177
178impl SelectorParser {
179    fn new(s: &str) -> SelectorParser {
180        SelectorParser {
181            chars: s.chars().collect(),
182            pos: 0,
183        }
184    }
185
186    fn parse(&mut self) -> Result<Vec<CompositeSelector>, String> {
187        let mut sels = vec![];
188        loop {
189            if self.cur().is_none() {
190                break;
191            }
192
193            let f1: OneSelector = if self.cur() == Some(':') {
194                OneSelector::Start
195            } else {
196                self.parse_one()?
197            };
198
199            let f2: Option<OneSelector> = if self.cur() == Some(':') {
200                self.bump();
201
202                let sel = if self.is_end_of_selector() {
203                    OneSelector::End
204                } else {
205                    self.parse_one()?
206                };
207
208                Some(sel)
209            } else {
210                None
211            };
212
213            if !self.is_end_of_selector() {
214                return Err(format!(
215                    "Expected end of field but got '{}' instead.",
216                    self.cur().unwrap()
217                ));
218            }
219
220            sels.push(match f2 {
221                Some(end) => CompositeSelector::Range(f1, end),
222                None => CompositeSelector::One(f1),
223            });
224
225            self.bump();
226        }
227
228        for sel in sels.iter_mut() {
229            sel.refine()?;
230        }
231
232        Ok(sels)
233    }
234
235    fn parse_one(&mut self) -> Result<OneSelector, String> {
236        let mut was_quoted = false;
237        let name = if self.cur() == Some('"') {
238            was_quoted = true;
239            self.bump();
240            self.parse_quoted_name()?
241        } else {
242            self.parse_name()?
243        };
244        Ok(if self.cur() == Some('[') {
245            let idx = self.parse_index()?;
246            OneSelector::IndexedName(name, Some(idx), was_quoted)
247        } else {
248            match name.parse() {
249                Err(_) => OneSelector::IndexedName(name, None, was_quoted),
250                Ok(idx) => OneSelector::Index(idx),
251            }
252        })
253    }
254
255    fn parse_name(&mut self) -> Result<String, String> {
256        let mut name = String::new();
257        loop {
258            if self.is_end_of_field() || self.cur() == Some('[') {
259                break;
260            }
261            name.push(self.cur().unwrap());
262            self.bump();
263        }
264        Ok(name)
265    }
266
267    fn parse_quoted_name(&mut self) -> Result<String, String> {
268        let mut name = String::new();
269        loop {
270            match self.cur() {
271                None => {
272                    return Err("Unclosed quote, missing closing \".".to_owned());
273                }
274                Some('"') => {
275                    self.bump();
276                    if self.cur() == Some('"') {
277                        self.bump();
278                        name.push('"');
279                        name.push('"');
280                        continue;
281                    }
282                    break;
283                }
284                Some(c) => {
285                    name.push(c);
286                    self.bump();
287                }
288            }
289        }
290        Ok(name)
291    }
292
293    fn parse_index(&mut self) -> Result<isize, String> {
294        assert_eq!(self.cur().unwrap(), '[');
295        self.bump();
296
297        let mut idx = String::new();
298        loop {
299            match self.cur() {
300                None => {
301                    return Err("Unclosed index bracket, missing closing ].".to_owned());
302                }
303                Some(']') => {
304                    self.bump();
305                    break;
306                }
307                Some(c) => {
308                    idx.push(c);
309                    self.bump();
310                }
311            }
312        }
313
314        idx.parse()
315            .map_err(|err| format!("Could not convert '{}' to an integer: {}", idx, err))
316    }
317
318    fn cur(&self) -> Option<char> {
319        self.chars.get(self.pos).cloned()
320    }
321
322    fn is_end_of_field(&self) -> bool {
323        match self.cur() {
324            None => true,
325            Some(c) => c == ',' || c == ':',
326        }
327    }
328
329    fn is_end_of_selector(&self) -> bool {
330        match self.cur() {
331            None => true,
332            Some(c) => c == ',',
333        }
334    }
335
336    fn bump(&mut self) {
337        if self.pos < self.chars.len() {
338            self.pos += 1;
339        }
340    }
341}
342
343#[derive(Clone)]
344enum CompositeSelector {
345    One(OneSelector),
346    Range(OneSelector, OneSelector),
347    GlobPrefix(String, Option<isize>),
348    GlobSuffix(String, Option<isize>),
349    GlobInner(String, String, Option<isize>),
350    All(Option<isize>),
351}
352
353impl CompositeSelector {
354    fn refine(&mut self) -> Result<(), String> {
355        match self {
356            Self::One(OneSelector::IndexedName(name, pos_opt, was_quoted)) => {
357                if *was_quoted {
358                    return Ok(());
359                }
360
361                let star_count = name.chars().filter(|c| *c == '*').count();
362
363                match star_count {
364                    0 => Ok(()),
365                    1 => {
366                        if name == "*" {
367                            *self = Self::All(*pos_opt);
368                        } else if name.starts_with('*') {
369                            *self = Self::GlobSuffix(
370                                name.trim_start_matches('*').to_string(),
371                                *pos_opt,
372                            );
373                        } else if name.ends_with('*') {
374                            *self =
375                                Self::GlobPrefix(name.trim_end_matches('*').to_string(), *pos_opt);
376                        } else {
377                            let pos = name
378                                .char_indices()
379                                .find_map(|(i, c)| if c == '*' { Some(i) } else { None })
380                                .unwrap();
381
382                            *self = Self::GlobInner(
383                                name[..pos].to_string(),
384                                name[pos + 1..].to_string(),
385                                *pos_opt,
386                            );
387                        }
388
389                        Ok(())
390                    }
391                    _ => Err(format!("'{}' contains more than one \"*\" wildcard", name)),
392                }
393            }
394            Self::Range(start, end) => {
395                if let OneSelector::IndexedName(name, _, false) = start {
396                    if name.contains("*") {
397                        return Err(
398                            "start of range cannot contain \"*\" wildcard unquoted".to_string()
399                        );
400                    }
401                }
402
403                if let OneSelector::IndexedName(name, _, false) = end {
404                    if name.contains("*") {
405                        return Err(
406                            "end of range cannot contain \"*\" wildcard unquoted".to_string()
407                        );
408                    }
409                }
410
411                Ok(())
412            }
413            _ => Ok(()),
414        }
415    }
416}
417
418#[derive(Clone)]
419enum OneSelector {
420    Start,
421    End,
422    Index(isize),
423    IndexedName(String, Option<isize>, bool),
424}
425
426impl CompositeSelector {
427    fn indices(&self, first_record: &[&[u8]], use_names: bool) -> Result<Vec<usize>, Error> {
428        struct Map<'s> {
429            inner: BTreeMap<&'s [u8], Vec<usize>>,
430        }
431
432        impl<'s> Map<'s> {
433            fn new(first_record: &'s [&[u8]]) -> Self {
434                let mut map = BTreeMap::new();
435
436                for (i, name) in first_record.iter().enumerate() {
437                    let list: &mut Vec<usize> = map.entry(*name).or_default();
438                    list.push(i);
439                }
440
441                Self { inner: map }
442            }
443
444            fn for_each<P, C>(&self, pos: isize, predicate: P, mut callback: C)
445            where
446                P: Fn(&[u8]) -> bool,
447                C: FnMut(usize),
448            {
449                for (name, indices) in self.inner.iter() {
450                    if !predicate(name) {
451                        continue;
452                    }
453
454                    let pos = if pos < 0 {
455                        indices.len() as isize + pos
456                    } else {
457                        pos
458                    };
459
460                    if pos < 0 {
461                        continue;
462                    }
463
464                    if let Some(i) = indices.get(pos as usize) {
465                        callback(*i);
466                    }
467                }
468            }
469        }
470
471        match *self {
472            CompositeSelector::All(pos_opt) => {
473                if let Some(pos) = pos_opt {
474                    if !use_names {
475                        return Err(Error::new(ErrorKind::SelectionError(format!(
476                            "Cannot use '*[{}]' in selection \
477                                        with --no-headers set.",
478                            pos
479                        ))));
480                    }
481
482                    let mut inds = vec![];
483                    let map = Map::new(first_record);
484
485                    map.for_each(pos, |_| true, |i| inds.push(i));
486
487                    if inds.is_empty() {
488                        return Err(Error::new(ErrorKind::SelectionError(format!(
489                            "'*[{}]' selected nothing.",
490                            pos
491                        ))));
492                    }
493
494                    Ok(inds)
495                } else {
496                    Ok((0..first_record.len()).collect())
497                }
498            }
499            CompositeSelector::One(ref sel) => sel.index(first_record, use_names).map(|i| vec![i]),
500            CompositeSelector::Range(ref sel1, ref sel2) => {
501                let i1 = sel1.index(first_record, use_names)?;
502                let i2 = sel2.index(first_record, use_names)?;
503                Ok(match i1.cmp(&i2) {
504                    Ordering::Equal => vec![i1],
505                    Ordering::Less => (i1..(i2 + 1)).collect(),
506                    Ordering::Greater => {
507                        let mut inds = vec![];
508                        let mut i = i1 + 1;
509                        while i > i2 {
510                            i -= 1;
511                            inds.push(i);
512                        }
513                        inds
514                    }
515                })
516            }
517            CompositeSelector::GlobPrefix(ref prefix, pos_opt) => {
518                if let Some(pos) = pos_opt {
519                    if !use_names {
520                        return Err(Error::new(ErrorKind::SelectionError(format!(
521                            "Cannot use prefix ('{}*[{}]') in selection \
522                                        with --no-headers set.",
523                            prefix, pos
524                        ))));
525                    }
526
527                    let mut inds = vec![];
528                    let map = Map::new(first_record);
529
530                    map.for_each(
531                        pos,
532                        |name| name.starts_with(prefix.as_bytes()),
533                        |i| inds.push(i),
534                    );
535
536                    if inds.is_empty() {
537                        return Err(Error::new(ErrorKind::SelectionError(format!(
538                            "Prefix '{}*[{}]' selected nothing.",
539                            prefix, pos
540                        ))));
541                    }
542
543                    Ok(inds)
544                } else {
545                    if !use_names {
546                        return Err(Error::new(ErrorKind::SelectionError(format!(
547                            "Cannot use prefix ('{}*') in selection \
548                                        with --no-headers set.",
549                            prefix
550                        ))));
551                    }
552
553                    let inds: Vec<usize> = first_record
554                        .iter()
555                        .enumerate()
556                        .filter_map(|(i, h)| {
557                            if h.starts_with(prefix.as_bytes()) {
558                                Some(i)
559                            } else {
560                                None
561                            }
562                        })
563                        .collect();
564
565                    if inds.is_empty() {
566                        return Err(Error::new(ErrorKind::SelectionError(format!(
567                            "Prefix '{}*' selected nothing.",
568                            prefix
569                        ))));
570                    }
571
572                    Ok(inds)
573                }
574            }
575            CompositeSelector::GlobSuffix(ref suffix, pos_opt) => {
576                if let Some(pos) = pos_opt {
577                    if !use_names {
578                        return Err(Error::new(ErrorKind::SelectionError(format!(
579                            "Cannot use suffix ('*{}[{}]') in selection \
580                                        with --no-headers set.",
581                            suffix, pos
582                        ))));
583                    }
584
585                    let mut inds = vec![];
586                    let map = Map::new(first_record);
587
588                    map.for_each(
589                        pos,
590                        |name| name.ends_with(suffix.as_bytes()),
591                        |i| inds.push(i),
592                    );
593
594                    if inds.is_empty() {
595                        return Err(Error::new(ErrorKind::SelectionError(format!(
596                            "Suffix '*{}[{}]' selected nothing.",
597                            suffix, pos
598                        ))));
599                    }
600
601                    Ok(inds)
602                } else {
603                    if !use_names {
604                        return Err(Error::new(ErrorKind::SelectionError(format!(
605                            "Cannot use suffix ('*{}') in selection \
606                                        with --no-headers set.",
607                            suffix
608                        ))));
609                    }
610
611                    let inds: Vec<usize> = first_record
612                        .iter()
613                        .enumerate()
614                        .filter_map(|(i, h)| {
615                            if h.ends_with(suffix.as_bytes()) {
616                                Some(i)
617                            } else {
618                                None
619                            }
620                        })
621                        .collect();
622
623                    if inds.is_empty() {
624                        return Err(Error::new(ErrorKind::SelectionError(format!(
625                            "Suffix '*{}' selected nothing.",
626                            suffix
627                        ))));
628                    }
629
630                    Ok(inds)
631                }
632            }
633            CompositeSelector::GlobInner(ref prefix, ref suffix, pos_opt) => {
634                if let Some(pos) = pos_opt {
635                    if !use_names {
636                        return Err(Error::new(ErrorKind::SelectionError(format!(
637                            "Cannot use inner wildcard ('{}*{}[{}]') in selection \
638                                        with --no-headers set.",
639                            prefix, suffix, pos
640                        ))));
641                    }
642
643                    let mut inds = vec![];
644                    let map = Map::new(first_record);
645
646                    map.for_each(
647                        pos,
648                        |name| {
649                            name.starts_with(prefix.as_bytes()) && name.ends_with(suffix.as_bytes())
650                        },
651                        |i| inds.push(i),
652                    );
653
654                    if inds.is_empty() {
655                        return Err(Error::new(ErrorKind::SelectionError(format!(
656                            "Inner wildcard '{}*{}[{}]' selected nothing.",
657                            prefix, suffix, pos
658                        ))));
659                    }
660
661                    Ok(inds)
662                } else {
663                    if !use_names {
664                        return Err(Error::new(ErrorKind::SelectionError(format!(
665                            "Cannot use inner wildcard ('{}*{}') in selection \
666                                        with --no-headers set.",
667                            prefix, suffix
668                        ))));
669                    }
670
671                    let inds: Vec<usize> = first_record
672                        .iter()
673                        .enumerate()
674                        .filter_map(|(i, h)| {
675                            if h.starts_with(prefix.as_bytes()) && h.ends_with(suffix.as_bytes()) {
676                                Some(i)
677                            } else {
678                                None
679                            }
680                        })
681                        .collect();
682
683                    if inds.is_empty() {
684                        return Err(Error::new(ErrorKind::SelectionError(format!(
685                            "Inner wildcard '{}*{}' selected nothing.",
686                            prefix, suffix
687                        ))));
688                    }
689
690                    Ok(inds)
691                }
692            }
693        }
694    }
695}
696
697impl OneSelector {
698    fn index(&self, first_record: &[&[u8]], use_names: bool) -> Result<usize, Error> {
699        match *self {
700            OneSelector::Start => Ok(0),
701            OneSelector::End => Ok(if first_record.is_empty() {
702                0
703            } else {
704                first_record.len() - 1
705            }),
706            OneSelector::Index(i) => {
707                if i < 0 {
708                    if i.unsigned_abs() > first_record.len() {
709                        Err(Error::new(ErrorKind::SelectionError(format!(
710                            "Column index {} is out of \
711                                 bounds. Index must be between -1 \
712                                 and -{}.",
713                            i,
714                            first_record.len()
715                        ))))
716                    } else {
717                        Ok(first_record.len() - i.unsigned_abs())
718                    }
719                } else {
720                    let i = i as usize;
721                    if i >= first_record.len() {
722                        Err(Error::new(ErrorKind::SelectionError(format!(
723                            "Column index {} is out of \
724                                 bounds. Index must be between 0 \
725                                 and {}.",
726                            i,
727                            first_record.len()
728                        ))))
729                    } else {
730                        Ok(i)
731                    }
732                }
733            }
734            OneSelector::IndexedName(ref s, sidx, _) => {
735                let sidx = sidx.unwrap_or(0);
736
737                if !use_names {
738                    return Err(Error::new(ErrorKind::SelectionError(format!(
739                        "Cannot use names ('{}') in selection \
740                                        with --no-headers set.",
741                        s
742                    ))));
743                }
744                let mut num_found = 0;
745
746                if sidx < 0 {
747                    for (i, field) in first_record.iter().enumerate().rev() {
748                        if field == &s.as_bytes() {
749                            if num_found == sidx.abs() - 1 {
750                                return Ok(i);
751                            }
752                            num_found += 1;
753                        }
754                    }
755                } else {
756                    for (i, field) in first_record.iter().enumerate() {
757                        if field == &s.as_bytes() {
758                            if num_found == sidx {
759                                return Ok(i);
760                            }
761                            num_found += 1;
762                        }
763                    }
764                }
765
766                if num_found == 0 {
767                    Err(Error::new(ErrorKind::SelectionError(format!(
768                        "'{}' does not exist \
769                                 as a named header in the given CSV \
770                                 data.",
771                        s
772                    ))))
773                } else if sidx < 0 {
774                    Err(Error::new(ErrorKind::SelectionError(format!(
775                        "index '{}' for '{}' is \
776                                     out of bounds. Must be between -{} and -1.",
777                        sidx, s, num_found
778                    ))))
779                } else {
780                    Err(Error::new(ErrorKind::SelectionError(format!(
781                        "index '{}' for name '{}' is \
782                                 out of bounds. Must be between 0 and {}.",
783                        sidx,
784                        s,
785                        num_found - 1
786                    ))))
787                }
788            }
789        }
790    }
791}
792
793impl fmt::Debug for CompositeSelector {
794    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
795        match *self {
796            CompositeSelector::All(pos_opt) => {
797                if let Some(pos) = pos_opt {
798                    write!(f, "All[{}]", pos)
799                } else {
800                    write!(f, "All")
801                }
802            }
803            CompositeSelector::One(ref sel) => sel.fmt(f),
804            CompositeSelector::Range(ref s, ref e) => write!(f, "Range({:?}, {:?})", s, e),
805            CompositeSelector::GlobPrefix(ref prefix, pos_opt) => write!(
806                f,
807                "Prefix({:?}){}",
808                prefix,
809                if let Some(pos) = pos_opt {
810                    format!("[{}]", pos)
811                } else {
812                    "".to_string()
813                }
814            ),
815            CompositeSelector::GlobSuffix(ref suffix, pos_opt) => write!(
816                f,
817                "Suffix({:?}){}",
818                suffix,
819                if let Some(pos) = pos_opt {
820                    format!("[{}]", pos)
821                } else {
822                    "".to_string()
823                }
824            ),
825            Self::GlobInner(ref prefix, ref suffix, pos_opt) => {
826                write!(
827                    f,
828                    "Inner({:?}, {:?}){}",
829                    prefix,
830                    suffix,
831                    if let Some(pos) = pos_opt {
832                        format!("[{}]", pos)
833                    } else {
834                        "".to_string()
835                    }
836                )
837            }
838        }
839    }
840}
841
842impl fmt::Debug for OneSelector {
843    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
844        match *self {
845            OneSelector::Start => write!(f, "Start"),
846            OneSelector::End => write!(f, "End"),
847            OneSelector::Index(idx) => write!(f, "Index({})", idx),
848            OneSelector::IndexedName(ref s, idx, _) => match idx {
849                None => write!(f, "IndexedName({})", s),
850                Some(i) => write!(f, "IndexedName({}[{}])", s, i),
851            },
852        }
853    }
854}
simd_csv/select/dsl.rs

simd_csv/select/
dsl.rs
