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

rustpython_vm/
suggestion.rs
