1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

use crate::{
    builtins::{PyStr, PyStrRef},
    exceptions::types::PyBaseExceptionRef,
    sliceable::SliceableSequenceOp,
    AsObject, Py, PyObjectRef, VirtualMachine,
};
use rustpython_common::str::levenshtein::{levenshtein_distance, MOVE_COST};
use std::iter::ExactSizeIterator;

const MAX_CANDIDATE_ITEMS: usize = 750;

fn calculate_suggestions<'a>(
    dir_iter: impl ExactSizeIterator<Item = &'a PyObjectRef>,
    name: &PyObjectRef,
) -> Option<PyStrRef> {
    if dir_iter.len() >= MAX_CANDIDATE_ITEMS {
        return None;
    }

    let mut suggestion: Option<&Py<PyStr>> = None;
    let mut suggestion_distance = usize::MAX;
    let name = name.downcast_ref::<PyStr>()?;

    for item in dir_iter {
        let item_name = item.downcast_ref::<PyStr>()?;
        if name.as_str() == item_name.as_str() {
            continue;
        }
        // No more than 1/3 of the characters should need changed
        let max_distance = usize::min(
            (name.len() + item_name.len() + 3) * MOVE_COST / 6,
            suggestion_distance - 1,
        );
        let current_distance =
            levenshtein_distance(name.as_str(), item_name.as_str(), max_distance);
        if current_distance > max_distance {
            continue;
        }
        if suggestion.is_none() || current_distance < suggestion_distance {
            suggestion = Some(item_name);
            suggestion_distance = current_distance;
        }
    }
    suggestion.map(|r| r.to_owned())
}

pub fn offer_suggestions(exc: &PyBaseExceptionRef, vm: &VirtualMachine) -> Option<PyStrRef> {
    if exc.class().is(vm.ctx.exceptions.attribute_error) {
        let name = exc.as_object().get_attr("name", vm).unwrap();
        let obj = exc.as_object().get_attr("obj", vm).unwrap();

        calculate_suggestions(vm.dir(Some(obj)).ok()?.borrow_vec().iter(), &name)
    } else if exc.class().is(vm.ctx.exceptions.name_error) {
        let name = exc.as_object().get_attr("name", vm).unwrap();
        let mut tb = exc.traceback()?;
        for traceback in tb.iter() {
            tb = traceback;
        }

        let varnames = tb.frame.code.clone().co_varnames(vm);
        if let Some(suggestions) = calculate_suggestions(varnames.iter(), &name) {
            return Some(suggestions);
        };

        let globals: Vec<_> = tb.frame.globals.as_object().try_to_value(vm).ok()?;
        if let Some(suggestions) = calculate_suggestions(globals.iter(), &name) {
            return Some(suggestions);
        };

        let builtins: Vec<_> = tb.frame.builtins.as_object().try_to_value(vm).ok()?;
        calculate_suggestions(builtins.iter(), &name)
    } else {
        None
    }
}