cpython/Python/suggestions.c

206 lines
6.1 KiB
C

#include "Python.h"
#include "pycore_code.h" // _PyCode_GetVarnames()
#include "pycore_frame.h"
#include "pycore_pyerrors.h" // export _Py_UTF8_Edit_Cost()
#include "pycore_runtime.h" // _Py_ID()
#define MAX_CANDIDATE_ITEMS 750
#define MAX_STRING_SIZE 40
#define MOVE_COST 2
#define CASE_COST 1
#define LEAST_FIVE_BITS(n) ((n) & 31)
static inline int
substitution_cost(char a, char b)
{
if (LEAST_FIVE_BITS(a) != LEAST_FIVE_BITS(b)) {
// Not the same, not a case flip.
return MOVE_COST;
}
if (a == b) {
return 0;
}
if ('A' <= a && a <= 'Z') {
a += ('a' - 'A');
}
if ('A' <= b && b <= 'Z') {
b += ('a' - 'A');
}
if (a == b) {
return CASE_COST;
}
return MOVE_COST;
}
/* Calculate the Levenshtein distance between string1 and string2 */
static Py_ssize_t
levenshtein_distance(const char *a, size_t a_size,
const char *b, size_t b_size,
size_t max_cost, size_t *buffer)
{
// Both strings are the same (by identity)
if (a == b) {
return 0;
}
// Trim away common affixes.
while (a_size && b_size && a[0] == b[0]) {
a++; a_size--;
b++; b_size--;
}
while (a_size && b_size && a[a_size-1] == b[b_size-1]) {
a_size--;
b_size--;
}
if (a_size == 0 || b_size == 0) {
return (a_size + b_size) * MOVE_COST;
}
if (a_size > MAX_STRING_SIZE || b_size > MAX_STRING_SIZE) {
return max_cost + 1;
}
// Prefer shorter buffer
if (b_size < a_size) {
const char *t = a; a = b; b = t;
size_t t_size = a_size; a_size = b_size; b_size = t_size;
}
// quick fail when a match is impossible.
if ((b_size - a_size) * MOVE_COST > max_cost) {
return max_cost + 1;
}
// Instead of producing the whole traditional len(a)-by-len(b)
// matrix, we can update just one row in place.
// Initialize the buffer row
size_t tmp = MOVE_COST;
for (size_t i = 0; i < a_size; i++) {
// cost from b[:0] to a[:i+1]
buffer[i] = tmp;
tmp += MOVE_COST;
}
size_t result = 0;
for (size_t b_index = 0; b_index < b_size; b_index++) {
char code = b[b_index];
// cost(b[:b_index], a[:0]) == b_index * MOVE_COST
size_t distance = result = b_index * MOVE_COST;
size_t minimum = SIZE_MAX;
for (size_t index = 0; index < a_size; index++) {
// cost(b[:b_index+1], a[:index+1]) = min(
// // 1) substitute
// cost(b[:b_index], a[:index])
// + substitution_cost(b[b_index], a[index]),
// // 2) delete from b
// cost(b[:b_index], a[:index+1]) + MOVE_COST,
// // 3) delete from a
// cost(b[:b_index+1], a[index]) + MOVE_COST
// )
// 1) Previous distance in this row is cost(b[:b_index], a[:index])
size_t substitute = distance + substitution_cost(code, a[index]);
// 2) cost(b[:b_index], a[:index+1]) from previous row
distance = buffer[index];
// 3) existing result is cost(b[:b_index+1], a[index])
size_t insert_delete = Py_MIN(result, distance) + MOVE_COST;
result = Py_MIN(insert_delete, substitute);
// cost(b[:b_index+1], a[:index+1])
buffer[index] = result;
if (result < minimum) {
minimum = result;
}
}
if (minimum > max_cost) {
// Everything in this row is too big, so bail early.
return max_cost + 1;
}
}
return result;
}
PyObject *
_Py_CalculateSuggestions(PyObject *dir,
PyObject *name)
{
assert(!PyErr_Occurred());
assert(PyList_CheckExact(dir));
Py_ssize_t dir_size = PyList_GET_SIZE(dir);
if (dir_size >= MAX_CANDIDATE_ITEMS) {
return NULL;
}
Py_ssize_t suggestion_distance = PY_SSIZE_T_MAX;
PyObject *suggestion = NULL;
Py_ssize_t name_size;
const char *name_str = PyUnicode_AsUTF8AndSize(name, &name_size);
if (name_str == NULL) {
return NULL;
}
size_t *buffer = PyMem_New(size_t, MAX_STRING_SIZE);
if (buffer == NULL) {
return PyErr_NoMemory();
}
for (Py_ssize_t i = 0; i < dir_size; ++i) {
PyObject *item = PyList_GET_ITEM(dir, i);
if (_PyUnicode_Equal(name, item)) {
continue;
}
Py_ssize_t item_size;
const char *item_str = PyUnicode_AsUTF8AndSize(item, &item_size);
if (item_str == NULL) {
PyMem_Free(buffer);
return NULL;
}
// No more than 1/3 of the involved characters should need changed.
Py_ssize_t max_distance = (name_size + item_size + 3) * MOVE_COST / 6;
// Don't take matches we've already beaten.
max_distance = Py_MIN(max_distance, suggestion_distance - 1);
Py_ssize_t current_distance =
levenshtein_distance(name_str, name_size, item_str,
item_size, max_distance, buffer);
if (current_distance > max_distance) {
continue;
}
if (!suggestion || current_distance < suggestion_distance) {
suggestion = item;
suggestion_distance = current_distance;
}
}
PyMem_Free(buffer);
return Py_XNewRef(suggestion);
}
Py_ssize_t
_Py_UTF8_Edit_Cost(PyObject *a, PyObject *b, Py_ssize_t max_cost)
{
assert(PyUnicode_Check(a) && PyUnicode_Check(b));
Py_ssize_t size_a, size_b;
const char *utf8_a = PyUnicode_AsUTF8AndSize(a, &size_a);
if (utf8_a == NULL) {
return -1;
}
const char *utf8_b = PyUnicode_AsUTF8AndSize(b, &size_b);
if (utf8_b == NULL) {
return -1;
}
if (max_cost == -1) {
max_cost = MOVE_COST * Py_MAX(size_a, size_b);
}
size_t *buffer = PyMem_New(size_t, MAX_STRING_SIZE);
if (buffer == NULL) {
PyErr_NoMemory();
return -1;
}
Py_ssize_t res = levenshtein_distance(utf8_a, size_a,
utf8_b, size_b, max_cost, buffer);
PyMem_Free(buffer);
return res;
}