dply/
completions.rs

1// Copyright (C) 2023 Vince Vasta
2// SPDX-License-Identifier: Apache-2.0
3use lru::LruCache;
4use polars::prelude::*;
5use regex::Regex;
6
7const MAX_ENTRIES: usize = 40;
8
9/// A completions LRU cache.
10pub struct Completions {
11    lru: LruCache<String, ()>,
12    name_re: regex::Regex,
13}
14
15impl Default for Completions {
16    fn default() -> Self {
17        Self {
18            lru: LruCache::unbounded(),
19            name_re: Regex::new(r"^[[:alpha:]](_|\d|[[:alpha:]])+$").unwrap(),
20        }
21    }
22}
23
24impl Completions {
25    /// Add entries to completions history.
26    pub fn add(&mut self, entries: &[PlSmallStr]) {
27        // Make sure these entries are in the cache irrespective of their size
28        // to handle the case where we have a dataframe with many columns.
29        if entries.len() > MAX_ENTRIES {
30            self.lru.clear();
31        } else if entries.len() + self.lru.len() > MAX_ENTRIES {
32            let to_remove = entries.len() + self.lru.len() - MAX_ENTRIES;
33            for _ in 0..to_remove {
34                self.lru.pop_lru();
35            }
36        }
37
38        for entry in entries {
39            self.add_entry(entry);
40        }
41    }
42
43    /// Returns an iterator to the completions.
44    pub fn iter(&self) -> impl Iterator<Item = &str> {
45        self.lru.iter().map(|(k, _)| k.as_str())
46    }
47
48    fn add_entry(&mut self, entry: &str) {
49        if self.lru.get(entry).is_none() {
50            // Add backticks to completions if entry is not a valid name.
51            let entry = if self.name_re.is_match(entry) {
52                entry.to_string()
53            } else {
54                format!("`{entry}`")
55            };
56
57            self.lru.put(entry, ());
58        }
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn completions() {
68        let mut completions = Completions::default();
69
70        // This insertion should have all columns, simulates a dataframe that has
71        // more columns than MAX_ENTRIES.
72        let entries = (0..MAX_ENTRIES + 10)
73            .map(|idx| PlSmallStr::from_string(format!("entry{idx}")))
74            .collect::<Vec<_>>();
75
76        completions.add(&entries);
77        assert_eq!(completions.iter().count(), entries.len());
78
79        // Inserts a subsets of columns smaller than MAX_ENTRIES.
80        let entries = (1000..1020)
81            .map(|idx| PlSmallStr::from_string(format!("entry{idx}")))
82            .collect::<Vec<_>>();
83        completions.add(&entries);
84
85        assert_eq!(completions.iter().count(), MAX_ENTRIES);
86
87        // Most recently used are first.
88        for (entry, (cached, _)) in entries.iter().rev().zip(completions.lru.iter()) {
89            assert_eq!(entry, cached);
90        }
91
92        let entries = (2000..2200)
93            .map(|idx| PlSmallStr::from_string(format!("entry{idx}")))
94            .collect::<Vec<_>>();
95        completions.add(&entries);
96        assert_eq!(completions.iter().count(), entries.len());
97
98        let entries = (3000..3100)
99            .map(|idx| PlSmallStr::from_string(format!("entry{idx}")))
100            .collect::<Vec<_>>();
101        completions.add(&entries);
102        assert_eq!(completions.iter().count(), entries.len());
103
104        for (entry, (cached, _)) in entries.iter().rev().zip(completions.lru.iter()) {
105            assert_eq!(entry, cached);
106        }
107    }
108}
dply/completions.rs

dply/
completions.rs
