rustpython_vm/
suggestion.rs

1use crate::{
2    builtins::{PyStr, PyStrRef},
3    exceptions::types::PyBaseExceptionRef,
4    sliceable::SliceableSequenceOp,
5    AsObject, Py, PyObjectRef, VirtualMachine,
6};
7use rustpython_common::str::levenshtein::{levenshtein_distance, MOVE_COST};
8use std::iter::ExactSizeIterator;
9
10const MAX_CANDIDATE_ITEMS: usize = 750;
11
12fn calculate_suggestions<'a>(
13    dir_iter: impl ExactSizeIterator<Item = &'a PyObjectRef>,
14    name: &PyObjectRef,
15) -> Option<PyStrRef> {
16    if dir_iter.len() >= MAX_CANDIDATE_ITEMS {
17        return None;
18    }
19
20    let mut suggestion: Option<&Py<PyStr>> = None;
21    let mut suggestion_distance = usize::MAX;
22    let name = name.downcast_ref::<PyStr>()?;
23
24    for item in dir_iter {
25        let item_name = item.downcast_ref::<PyStr>()?;
26        if name.as_str() == item_name.as_str() {
27            continue;
28        }
29        // No more than 1/3 of the characters should need changed
30        let max_distance = usize::min(
31            (name.len() + item_name.len() + 3) * MOVE_COST / 6,
32            suggestion_distance - 1,
33        );
34        let current_distance =
35            levenshtein_distance(name.as_str(), item_name.as_str(), max_distance);
36        if current_distance > max_distance {
37            continue;
38        }
39        if suggestion.is_none() || current_distance < suggestion_distance {
40            suggestion = Some(item_name);
41            suggestion_distance = current_distance;
42        }
43    }
44    suggestion.map(|r| r.to_owned())
45}
46
47pub fn offer_suggestions(exc: &PyBaseExceptionRef, vm: &VirtualMachine) -> Option<PyStrRef> {
48    if exc.class().is(vm.ctx.exceptions.attribute_error) {
49        let name = exc.as_object().get_attr("name", vm).unwrap();
50        let obj = exc.as_object().get_attr("obj", vm).unwrap();
51
52        calculate_suggestions(vm.dir(Some(obj)).ok()?.borrow_vec().iter(), &name)
53    } else if exc.class().is(vm.ctx.exceptions.name_error) {
54        let name = exc.as_object().get_attr("name", vm).unwrap();
55        let mut tb = exc.traceback()?;
56        for traceback in tb.iter() {
57            tb = traceback;
58        }
59
60        let varnames = tb.frame.code.clone().co_varnames(vm);
61        if let Some(suggestions) = calculate_suggestions(varnames.iter(), &name) {
62            return Some(suggestions);
63        };
64
65        let globals: Vec<_> = tb.frame.globals.as_object().try_to_value(vm).ok()?;
66        if let Some(suggestions) = calculate_suggestions(globals.iter(), &name) {
67            return Some(suggestions);
68        };
69
70        let builtins: Vec<_> = tb.frame.builtins.as_object().try_to_value(vm).ok()?;
71        calculate_suggestions(builtins.iter(), &name)
72    } else {
73        None
74    }
75}
rustpython_vm/suggestion.rs

rustpython_vm/
suggestion.rs
