rustpython_vm/
suggestion.rs

1//! This module provides functionality to suggest similar names for attributes or variables.
2//! This is used during tracebacks.
3
4use crate::{
5    AsObject, Py, PyObject, PyObjectRef, VirtualMachine,
6    builtins::{PyStr, PyStrRef},
7    exceptions::types::PyBaseException,
8    sliceable::SliceableSequenceOp,
9};
10use core::iter::ExactSizeIterator;
11use rustpython_common::str::levenshtein::{MOVE_COST, levenshtein_distance};
12
13const MAX_CANDIDATE_ITEMS: usize = 750;
14
15pub fn calculate_suggestions<'a>(
16    dir_iter: impl ExactSizeIterator<Item = &'a PyObjectRef>,
17    name: &PyObject,
18) -> Option<PyStrRef> {
19    if dir_iter.len() >= MAX_CANDIDATE_ITEMS {
20        return None;
21    }
22
23    let mut suggestion: Option<&Py<PyStr>> = None;
24    let mut suggestion_distance = usize::MAX;
25    let name = name.downcast_ref::<PyStr>()?;
26
27    for item in dir_iter {
28        let item_name = item.downcast_ref::<PyStr>()?;
29        if name.as_bytes() == item_name.as_bytes() {
30            continue;
31        }
32        // No more than 1/3 of the characters should need changed
33        let max_distance = usize::min(
34            (name.len() + item_name.len() + 3) * MOVE_COST / 6,
35            suggestion_distance - 1,
36        );
37        let current_distance =
38            levenshtein_distance(name.as_bytes(), item_name.as_bytes(), max_distance);
39        if current_distance > max_distance {
40            continue;
41        }
42        if suggestion.is_none() || current_distance < suggestion_distance {
43            suggestion = Some(item_name);
44            suggestion_distance = current_distance;
45        }
46    }
47    suggestion.map(|r| r.to_owned())
48}
49
50pub fn offer_suggestions(exc: &Py<PyBaseException>, vm: &VirtualMachine) -> Option<PyStrRef> {
51    if exc
52        .class()
53        .fast_issubclass(vm.ctx.exceptions.attribute_error)
54    {
55        let name = exc.as_object().get_attr("name", vm).ok()?;
56        if vm.is_none(&name) {
57            return None;
58        }
59        let obj = exc.as_object().get_attr("obj", vm).ok()?;
60        if vm.is_none(&obj) {
61            return None;
62        }
63
64        calculate_suggestions(vm.dir(Some(obj)).ok()?.borrow_vec().iter(), &name)
65    } else if exc.class().fast_issubclass(vm.ctx.exceptions.name_error) {
66        let name = exc.as_object().get_attr("name", vm).ok()?;
67        if vm.is_none(&name) {
68            return None;
69        }
70        let tb = exc.__traceback__()?;
71        let tb = tb.iter().last().unwrap_or(tb);
72
73        let varnames = tb.frame.code.clone().co_varnames(vm);
74        if let Some(suggestions) = calculate_suggestions(varnames.iter(), &name) {
75            return Some(suggestions);
76        };
77
78        let globals: Vec<_> = tb.frame.globals.as_object().try_to_value(vm).ok()?;
79        if let Some(suggestions) = calculate_suggestions(globals.iter(), &name) {
80            return Some(suggestions);
81        };
82
83        let builtins: Vec<_> = tb.frame.builtins.try_to_value(vm).ok()?;
84        calculate_suggestions(builtins.iter(), &name)
85    } else if exc.class().fast_issubclass(vm.ctx.exceptions.import_error) {
86        let mod_name = exc.as_object().get_attr("name", vm).ok()?;
87        let wrong_name = exc.as_object().get_attr("name_from", vm).ok()?;
88        let mod_name_str = mod_name.downcast_ref::<PyStr>()?;
89
90        // Look up the module in sys.modules
91        let sys_modules = vm.sys_module.get_attr("modules", vm).ok()?;
92        let module = sys_modules.get_item(mod_name_str, vm).ok()?;
93
94        calculate_suggestions(vm.dir(Some(module)).ok()?.borrow_vec().iter(), &wrong_name)
95    } else {
96        None
97    }
98}
rustpython_vm/suggestion.rs

rustpython_vm/
suggestion.rs
