diff --git a/src/shell/helper.rs b/src/shell/helper.rs index 493dc57846..8d0267b120 100644 --- a/src/shell/helper.rs +++ b/src/shell/helper.rs @@ -1,4 +1,5 @@ #![cfg_attr(target_arch = "wasm32", allow(dead_code))] + use rustpython_vm::{ AsObject, PyResult, TryFromObject, VirtualMachine, builtins::{PyDictRef, PyStrRef}, @@ -11,138 +12,113 @@ pub struct ShellHelper<'vm> { globals: PyDictRef, } -fn reverse_string(s: &mut String) { - let rev = s.chars().rev().collect(); - *s = rev; -} - -fn split_idents_on_dot(line: &str) -> Option<(usize, Vec)> { - let mut words = vec![String::new()]; - let mut startpos = 0; - for (i, c) in line.chars().rev().enumerate() { - match c { - '.' => { - // check for a double dot - if i != 0 && words.last().is_some_and(|s| s.is_empty()) { - return None; - } - reverse_string(words.last_mut().unwrap()); - if words.len() == 1 { - startpos = line.len() - i; - } - words.push(String::new()); - } - c if c.is_alphanumeric() || c == '_' => words.last_mut().unwrap().push(c), - _ => { - if words.len() == 1 { - if words.last().unwrap().is_empty() { - return None; - } - startpos = line.len() - i; - } - break; - } - } - } - if words == [String::new()] { - return None; - } - reverse_string(words.last_mut().unwrap()); - words.reverse(); - - Some((startpos, words)) -} - impl<'vm> ShellHelper<'vm> { pub const fn new(vm: &'vm VirtualMachine, globals: PyDictRef) -> Self { - ShellHelper { vm, globals } + Self { vm, globals } } fn get_available_completions<'w>( &self, words: &'w [String], ) -> Option<(&'w str, impl Iterator> + 'vm)> { - // the very first word and then all the ones after the dot - let (first, rest) = words.split_first().unwrap(); + let (first, rest) = words.split_first()?; - let str_iter_method = |obj, name| { + let get_str_iter = |obj, name| { let iter = self.vm.call_special_method(obj, name, ())?; ArgIterable::::try_from_object(self.vm, iter)?.iter(self.vm) }; - let (word_start, iter1, iter2) = if let Some((last, parents)) = rest.split_last() { - // we need to get an attribute based off of the dir() of an object - - // last: the last word, could be empty if it ends with a dot - // parents: the words before the dot - - let mut current = self.globals.get_item_opt(first.as_str(), self.vm).ok()??; - + if let Some((last, parents)) = rest.split_last() { + let mut current = self.globals.get_item_opt(first, self.vm).ok()??; for attr in parents { - let attr = self.vm.ctx.new_str(attr.as_str()); - current = current.get_attr(&attr, self.vm).ok()?; + current = current.get_attr(self.vm.ctx.new_str(attr), self.vm).ok()?; } - - let current_iter = str_iter_method(¤t, identifier!(self.vm, __dir__)).ok()?; - - (last, current_iter, None) + let iter = get_str_iter(¤t, identifier!(self.vm, __dir__)).ok()?; + Some((last, iter)) } else { - // we need to get a variable based off of globals/builtins - - let globals = - str_iter_method(self.globals.as_object(), identifier!(self.vm, keys)).ok()?; - let builtins = - str_iter_method(self.vm.builtins.as_object(), identifier!(self.vm, __dir__)) - .ok()?; - (first, globals, Some(builtins)) - }; - Some((word_start, iter1.chain(iter2.into_iter().flatten()))) + let globals = get_str_iter(self.globals.as_object(), identifier!(self.vm, keys)).ok()?; + let builtins = get_str_iter(self.vm.builtins.as_object(), identifier!(self.vm, __dir__)).ok()?; + Some((first, globals.chain(builtins))) + } } fn complete_opt(&self, line: &str) -> Option<(usize, Vec)> { let (startpos, words) = split_idents_on_dot(line)?; + let (prefix, completions) = self.get_available_completions(&words)?; - let (word_start, iter) = self.get_available_completions(&words)?; - - let all_completions = iter - .filter(|res| { - res.as_ref() - .ok() - .is_none_or(|s| s.as_str().starts_with(word_start)) - }) - .collect::, _>>() - .ok()?; - let mut completions = if word_start.starts_with('_') { - // if they're already looking for something starting with a '_', just give - // them all the completions - all_completions + let completions = completions + .filter_map(Result::ok) + .filter(|s| prefix.is_empty() || s.as_str().starts_with(prefix)) + .collect::>(); + + let filtered = if prefix.starts_with('_') { + completions } else { - // only the completions that don't start with a '_' - let no_underscore = all_completions + let no_underscore: Vec<_> = completions .iter() - .filter(|&s| !s.as_str().starts_with('_')) + .filter(|s| !s.as_str().starts_with('_')) .cloned() - .collect::>(); - - // if there are only completions that start with a '_', give them all of the - // completions, otherwise only the ones that don't start with '_' + .collect(); if no_underscore.is_empty() { - all_completions + completions } else { no_underscore } }; - // sort the completions alphabetically - completions.sort_by(|a, b| std::cmp::Ord::cmp(a.as_str(), b.as_str())); + let mut result = filtered + .into_iter() + .map(|s| s.as_str().to_owned()) + .collect::>(); + result.sort_unstable(); - Some(( - startpos, - completions - .into_iter() - .map(|s| s.as_str().to_owned()) - .collect(), - )) + Some((startpos, result)) + } +} + +fn split_idents_on_dot(line: &str) -> Option<(usize, Vec)> { + let mut idents = vec![String::new()]; + let mut startpos = 0; + let mut rev_chars = line.chars().rev().enumerate(); + + while let Some((i, c)) = rev_chars.next() { + match c { + '.' => { + if idents.last().map_or(false, |s| s.is_empty()) && i != 0 { + return None; // Double dot + } + idents.last_mut().map(|s| reverse_string(s)); + if idents.len() == 1 { + startpos = line.len() - i; + } + idents.push(String::new()); + } + c if c.is_alphanumeric() || c == '_' => { + idents.last_mut()?.push(c); + } + _ => { + if idents.len() == 1 && idents.last().map_or(true, |s| s.is_empty()) { + return None; + } + startpos = line.len() - i; + break; + } + } + } + + if idents == [String::new()] { + return None; + } + + idents.last_mut().map(|s| reverse_string(s)); + idents.reverse(); + Some((startpos, idents)) +} + +fn reverse_string(s: &mut String) { + unsafe { + let bytes = s.as_bytes_mut(); + bytes.reverse(); } } @@ -152,6 +128,7 @@ cfg_if::cfg_if! { completion::Completer, highlight::Highlighter, hint::Hinter, validate::Validator, Context, Helper, }; + impl Completer for ShellHelper<'_> { type Candidate = String; @@ -161,10 +138,7 @@ cfg_if::cfg_if! { pos: usize, _ctx: &Context, ) -> rustyline::Result<(usize, Vec)> { - Ok(self - .complete_opt(&line[0..pos]) - // as far as I can tell, there's no better way to do both completion - // and indentation (or even just indentation) + Ok(self.complete_opt(&line[..pos]) .unwrap_or_else(|| (pos, vec!["\t".to_owned()]))) } }