3
3
4
4
#include "pycore_pyerrors.h"
5
5
#include "pycore_code.h" // _PyCode_GetVarnames()
6
+ #include "stdlib_module_names.h" // _Py_stdlib_module_names
6
7
7
8
#define MAX_CANDIDATE_ITEMS 750
8
9
#define MAX_STRING_SIZE 40
@@ -175,7 +176,7 @@ calculate_suggestions(PyObject *dir,
175
176
}
176
177
177
178
static PyObject *
178
- offer_suggestions_for_attribute_error (PyAttributeErrorObject * exc )
179
+ get_suggestions_for_attribute_error (PyAttributeErrorObject * exc )
179
180
{
180
181
PyObject * name = exc -> name ; // borrowed reference
181
182
PyObject * obj = exc -> obj ; // borrowed reference
@@ -195,35 +196,25 @@ offer_suggestions_for_attribute_error(PyAttributeErrorObject *exc)
195
196
return suggestions ;
196
197
}
197
198
198
-
199
199
static PyObject *
200
- offer_suggestions_for_name_error ( PyNameErrorObject * exc )
200
+ offer_suggestions_for_attribute_error ( PyAttributeErrorObject * exc )
201
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
- ) {
202
+ PyObject * suggestion = get_suggestions_for_attribute_error (exc );
203
+ if (suggestion == NULL ) {
209
204
return NULL ;
210
205
}
206
+ // Add a trailer ". Did you mean: (...)?"
207
+ PyObject * result = PyUnicode_FromFormat (". Did you mean: %R?" , suggestion );
208
+ Py_DECREF (suggestion );
209
+ return result ;
210
+ }
211
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 );
212
+ static PyObject *
213
+ get_suggestions_for_name_error (PyObject * name , PyFrameObject * frame )
214
+ {
225
215
PyCodeObject * code = PyFrame_GetCode (frame );
226
216
assert (code != NULL && code -> co_localsplusnames != NULL );
217
+
227
218
PyObject * varnames = _PyCode_GetVarnames (code );
228
219
if (varnames == NULL ) {
229
220
return NULL ;
@@ -261,6 +252,66 @@ offer_suggestions_for_name_error(PyNameErrorObject *exc)
261
252
return suggestions ;
262
253
}
263
254
255
+ static bool
256
+ is_name_stdlib_module (PyObject * name )
257
+ {
258
+ const char * the_name = PyUnicode_AsUTF8 (name );
259
+ Py_ssize_t len = Py_ARRAY_LENGTH (_Py_stdlib_module_names );
260
+ for (Py_ssize_t i = 0 ; i < len ; i ++ ) {
261
+ if (strcmp (the_name , _Py_stdlib_module_names [i ]) == 0 ) {
262
+ return 1 ;
263
+ }
264
+ }
265
+ return 0 ;
266
+ }
267
+
268
+ static PyObject *
269
+ offer_suggestions_for_name_error (PyNameErrorObject * exc )
270
+ {
271
+ PyObject * name = exc -> name ; // borrowed reference
272
+ PyTracebackObject * traceback = (PyTracebackObject * ) exc -> traceback ; // borrowed reference
273
+ // Abort if we don't have a variable name or we have an invalid one
274
+ // or if we don't have a traceback to work with
275
+ if (name == NULL || !PyUnicode_CheckExact (name ) ||
276
+ traceback == NULL || !Py_IS_TYPE (traceback , & PyTraceBack_Type )
277
+ ) {
278
+ return NULL ;
279
+ }
280
+
281
+ // Move to the traceback of the exception
282
+ while (1 ) {
283
+ PyTracebackObject * next = traceback -> tb_next ;
284
+ if (next == NULL || !Py_IS_TYPE (next , & PyTraceBack_Type )) {
285
+ break ;
286
+ }
287
+ else {
288
+ traceback = next ;
289
+ }
290
+ }
291
+
292
+ PyFrameObject * frame = traceback -> tb_frame ;
293
+ assert (frame != NULL );
294
+
295
+ PyObject * suggestion = get_suggestions_for_name_error (name , frame );
296
+ bool is_stdlib_module = is_name_stdlib_module (name );
297
+
298
+ if (suggestion == NULL && !is_stdlib_module ) {
299
+ return NULL ;
300
+ }
301
+
302
+ // Add a trailer ". Did you mean: (...)?"
303
+ PyObject * result = NULL ;
304
+ if (!is_stdlib_module ) {
305
+ result = PyUnicode_FromFormat (". Did you mean: %R?" , suggestion );
306
+ } else if (suggestion == NULL ) {
307
+ result = PyUnicode_FromFormat (". Did you forget to import %R?" , name );
308
+ } else {
309
+ result = PyUnicode_FromFormat (". Did you mean: %R? Or did you forget to import %R?" , suggestion , name );
310
+ }
311
+ Py_XDECREF (suggestion );
312
+ return result ;
313
+ }
314
+
264
315
// Offer suggestions for a given exception. Returns a python string object containing the
265
316
// suggestions. This function returns NULL if no suggestion was found or if an exception happened,
266
317
// users must call PyErr_Occurred() to disambiguate.
0 commit comments