1 #include "Python.h"
2 #include "pycore_frame.h"
3 
4 #include "pycore_pyerrors.h"
5 #include "pycore_code.h"        // _PyCode_GetVarnames()
6 
7 #define MAX_CANDIDATE_ITEMS 750
8 #define MAX_STRING_SIZE 40
9 
10 #define MOVE_COST 2
11 #define CASE_COST 1
12 
13 #define LEAST_FIVE_BITS(n) ((n) & 31)
14 
15 static inline int
substitution_cost(char a,char b)16 substitution_cost(char a, char b)
17 {
18     if (LEAST_FIVE_BITS(a) != LEAST_FIVE_BITS(b)) {
19         // Not the same, not a case flip.
20         return MOVE_COST;
21     }
22     if (a == b) {
23         return 0;
24     }
25     if ('A' <= a && a <= 'Z') {
26         a += ('a' - 'A');
27     }
28     if ('A' <= b && b <= 'Z') {
29         b += ('a' - 'A');
30     }
31     if (a == b) {
32         return CASE_COST;
33     }
34     return MOVE_COST;
35 }
36 
37 /* Calculate the Levenshtein distance between string1 and string2 */
38 static Py_ssize_t
levenshtein_distance(const char * a,size_t a_size,const char * b,size_t b_size,size_t max_cost)39 levenshtein_distance(const char *a, size_t a_size,
40                      const char *b, size_t b_size,
41                      size_t max_cost)
42 {
43     static size_t buffer[MAX_STRING_SIZE];
44 
45     // Both strings are the same (by identity)
46     if (a == b) {
47         return 0;
48     }
49 
50     // Trim away common affixes.
51     while (a_size && b_size && a[0] == b[0]) {
52         a++; a_size--;
53         b++; b_size--;
54     }
55     while (a_size && b_size && a[a_size-1] == b[b_size-1]) {
56         a_size--;
57         b_size--;
58     }
59     if (a_size == 0 || b_size == 0) {
60         return (a_size + b_size) * MOVE_COST;
61     }
62     if (a_size > MAX_STRING_SIZE || b_size > MAX_STRING_SIZE) {
63         return max_cost + 1;
64     }
65 
66     // Prefer shorter buffer
67     if (b_size < a_size) {
68         const char *t = a; a = b; b = t;
69         size_t t_size = a_size; a_size = b_size; b_size = t_size;
70     }
71 
72     // quick fail when a match is impossible.
73     if ((b_size - a_size) * MOVE_COST > max_cost) {
74         return max_cost + 1;
75     }
76 
77     // Instead of producing the whole traditional len(a)-by-len(b)
78     // matrix, we can update just one row in place.
79     // Initialize the buffer row
80     size_t tmp = MOVE_COST;
81     for (size_t i = 0; i < a_size; i++) {
82         // cost from b[:0] to a[:i+1]
83         buffer[i] = tmp;
84         tmp += MOVE_COST;
85     }
86 
87     size_t result = 0;
88     for (size_t b_index = 0; b_index < b_size; b_index++) {
89         char code = b[b_index];
90         // cost(b[:b_index], a[:0]) == b_index * MOVE_COST
91         size_t distance = result = b_index * MOVE_COST;
92         size_t minimum = SIZE_MAX;
93         for (size_t index = 0; index < a_size; index++) {
94 
95             // cost(b[:b_index+1], a[:index+1]) = min(
96             //     // 1) substitute
97             //     cost(b[:b_index], a[:index])
98             //         + substitution_cost(b[b_index], a[index]),
99             //     // 2) delete from b
100             //     cost(b[:b_index], a[:index+1]) + MOVE_COST,
101             //     // 3) delete from a
102             //     cost(b[:b_index+1], a[index]) + MOVE_COST
103             // )
104 
105             // 1) Previous distance in this row is cost(b[:b_index], a[:index])
106             size_t substitute = distance + substitution_cost(code, a[index]);
107             // 2) cost(b[:b_index], a[:index+1]) from previous row
108             distance = buffer[index];
109             // 3) existing result is cost(b[:b_index+1], a[index])
110 
111             size_t insert_delete = Py_MIN(result, distance) + MOVE_COST;
112             result = Py_MIN(insert_delete, substitute);
113 
114             // cost(b[:b_index+1], a[:index+1])
115             buffer[index] = result;
116             if (result < minimum) {
117                 minimum = result;
118             }
119         }
120         if (minimum > max_cost) {
121             // Everything in this row is too big, so bail early.
122             return max_cost + 1;
123         }
124     }
125     return result;
126 }
127 
128 static inline PyObject *
calculate_suggestions(PyObject * dir,PyObject * name)129 calculate_suggestions(PyObject *dir,
130                       PyObject *name)
131 {
132     assert(!PyErr_Occurred());
133     assert(PyList_CheckExact(dir));
134 
135     Py_ssize_t dir_size = PyList_GET_SIZE(dir);
136     if (dir_size >= MAX_CANDIDATE_ITEMS) {
137         return NULL;
138     }
139 
140     Py_ssize_t suggestion_distance = PY_SSIZE_T_MAX;
141     PyObject *suggestion = NULL;
142     Py_ssize_t name_size;
143     const char *name_str = PyUnicode_AsUTF8AndSize(name, &name_size);
144     if (name_str == NULL) {
145         return NULL;
146     }
147 
148     for (int i = 0; i < dir_size; ++i) {
149         PyObject *item = PyList_GET_ITEM(dir, i);
150         Py_ssize_t item_size;
151         const char *item_str = PyUnicode_AsUTF8AndSize(item, &item_size);
152         if (item_str == NULL) {
153             return NULL;
154         }
155         if (PyUnicode_CompareWithASCIIString(name, item_str) == 0) {
156             continue;
157         }
158         // No more than 1/3 of the involved characters should need changed.
159         Py_ssize_t max_distance = (name_size + item_size + 3) * MOVE_COST / 6;
160         // Don't take matches we've already beaten.
161         max_distance = Py_MIN(max_distance, suggestion_distance - 1);
162         Py_ssize_t current_distance =
163             levenshtein_distance(name_str, name_size,
164                                  item_str, item_size, max_distance);
165         if (current_distance > max_distance) {
166             continue;
167         }
168         if (!suggestion || current_distance < suggestion_distance) {
169             suggestion = item;
170             suggestion_distance = current_distance;
171         }
172     }
173     Py_XINCREF(suggestion);
174     return suggestion;
175 }
176 
177 static PyObject *
offer_suggestions_for_attribute_error(PyAttributeErrorObject * exc)178 offer_suggestions_for_attribute_error(PyAttributeErrorObject *exc)
179 {
180     PyObject *name = exc->name; // borrowed reference
181     PyObject *obj = exc->obj; // borrowed reference
182 
183     // Abort if we don't have an attribute name or we have an invalid one
184     if (name == NULL || obj == NULL || !PyUnicode_CheckExact(name)) {
185         return NULL;
186     }
187 
188     PyObject *dir = PyObject_Dir(obj);
189     if (dir == NULL) {
190         return NULL;
191     }
192 
193     PyObject *suggestions = calculate_suggestions(dir, name);
194     Py_DECREF(dir);
195     return suggestions;
196 }
197 
198 
199 static PyObject *
offer_suggestions_for_name_error(PyNameErrorObject * exc)200 offer_suggestions_for_name_error(PyNameErrorObject *exc)
201 {
202     PyObject *name = exc->name; // borrowed reference
203     PyTracebackObject *traceback = (PyTracebackObject *) exc->traceback; // borrowed reference
204     // Abort if we don't have a variable name or we have an invalid one
205     // or if we don't have a traceback to work with
206     if (name == NULL || !PyUnicode_CheckExact(name) ||
207         traceback == NULL || !Py_IS_TYPE(traceback, &PyTraceBack_Type)
208     ) {
209         return NULL;
210     }
211 
212     // Move to the traceback of the exception
213     while (1) {
214         PyTracebackObject *next = traceback->tb_next;
215         if (next == NULL || !Py_IS_TYPE(next, &PyTraceBack_Type)) {
216             break;
217         }
218         else {
219             traceback = next;
220         }
221     }
222 
223     PyFrameObject *frame = traceback->tb_frame;
224     assert(frame != NULL);
225     PyCodeObject *code = PyFrame_GetCode(frame);
226     assert(code != NULL && code->co_localsplusnames != NULL);
227     PyObject *varnames = _PyCode_GetVarnames(code);
228     if (varnames == NULL) {
229         return NULL;
230     }
231     PyObject *dir = PySequence_List(varnames);
232     Py_DECREF(varnames);
233     Py_DECREF(code);
234     if (dir == NULL) {
235         return NULL;
236     }
237 
238     PyObject *suggestions = calculate_suggestions(dir, name);
239     Py_DECREF(dir);
240     if (suggestions != NULL) {
241         return suggestions;
242     }
243 
244     dir = PySequence_List(frame->f_frame->f_globals);
245     if (dir == NULL) {
246         return NULL;
247     }
248     suggestions = calculate_suggestions(dir, name);
249     Py_DECREF(dir);
250     if (suggestions != NULL) {
251         return suggestions;
252     }
253 
254     dir = PySequence_List(frame->f_frame->f_builtins);
255     if (dir == NULL) {
256         return NULL;
257     }
258     suggestions = calculate_suggestions(dir, name);
259     Py_DECREF(dir);
260 
261     return suggestions;
262 }
263 
264 // Offer suggestions for a given exception. Returns a python string object containing the
265 // suggestions. This function returns NULL if no suggestion was found or if an exception happened,
266 // users must call PyErr_Occurred() to disambiguate.
267 PyObject *
_Py_Offer_Suggestions(PyObject * exception)268 _Py_Offer_Suggestions(PyObject *exception)
269 {
270     PyObject *result = NULL;
271     assert(!PyErr_Occurred());
272     if (Py_IS_TYPE(exception, (PyTypeObject*)PyExc_AttributeError)) {
273         result = offer_suggestions_for_attribute_error((PyAttributeErrorObject *) exception);
274     } else if (Py_IS_TYPE(exception, (PyTypeObject*)PyExc_NameError)) {
275         result = offer_suggestions_for_name_error((PyNameErrorObject *) exception);
276     }
277     return result;
278 }
279 
280 Py_ssize_t
_Py_UTF8_Edit_Cost(PyObject * a,PyObject * b,Py_ssize_t max_cost)281 _Py_UTF8_Edit_Cost(PyObject *a, PyObject *b, Py_ssize_t max_cost)
282 {
283     assert(PyUnicode_Check(a) && PyUnicode_Check(b));
284     Py_ssize_t size_a, size_b;
285     const char *utf8_a = PyUnicode_AsUTF8AndSize(a, &size_a);
286     if (utf8_a == NULL) {
287         return -1;
288     }
289     const char *utf8_b = PyUnicode_AsUTF8AndSize(b, &size_b);
290     if (utf8_b == NULL) {
291         return -1;
292     }
293     if (max_cost == -1) {
294         max_cost = MOVE_COST * Py_MAX(size_a, size_b);
295     }
296     return levenshtein_distance(utf8_a, size_a, utf8_b, size_b, max_cost);
297 }
298 
299