Skip to content

Commit 66f9385

Browse files
committed
pythongh-119118: Fix performance regression in tokenize module
- Cache line object to avoid creating a Unicode object for all of the tokens in the same line. - Speed up byte offset to column offset conversion by using the smallest buffer possible to measure the difference.
1 parent 65de194 commit 66f9385

1 file changed

Lines changed: 44 additions & 4 deletions

File tree

‎Python/Python-tokenize.c‎

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,12 @@ typedef struct
3232
{
3333
PyObject_HEAD struct tok_state *tok;
3434
int done;
35+
36+
/* Needed to cache line for performance */
37+
const char *last_line_start;
38+
Py_ssize_t last_line_size;
39+
PyObject *last_line;
40+
Py_ssize_t byte_col_offset_diff;
3541
} tokenizeriterobject;
3642

3743
/*[clinic input]
@@ -68,6 +74,12 @@ tokenizeriter_new_impl(PyTypeObject *type, PyObject *readline,
6874
self->tok->tok_extra_tokens = 1;
6975
}
7076
self->done = 0;
77+
78+
self->last_line_start = NULL;
79+
self->last_line_size = -2;
80+
self->last_line = NULL;
81+
self->byte_col_offset_diff = 0;
82+
7183
return (PyObject *)self;
7284
}
7385

@@ -210,7 +222,20 @@ tokenizeriter_next(tokenizeriterobject *it)
210222
if (size >= 1 && it->tok->implicit_newline) {
211223
size -= 1;
212224
}
213-
line = PyUnicode_DecodeUTF8(line_start, size, "replace");
225+
226+
if (line_start != it->last_line_start || size != it->last_line_size) {
227+
// Line has changed since last token, so we fetch the new line and cache it
228+
// in the iter object.
229+
Py_XDECREF(it->last_line);
230+
line = PyUnicode_DecodeUTF8(line_start, size, "replace");
231+
it->last_line = line;
232+
it->last_line_start = line_start;
233+
it->last_line_size = size;
234+
it->byte_col_offset_diff = 0;
235+
} else {
236+
// Line hasn't changed so we reuse the cached one.
237+
line = it->last_line;
238+
}
214239
}
215240
if (line == NULL) {
216241
Py_DECREF(str);
@@ -222,10 +247,25 @@ tokenizeriter_next(tokenizeriterobject *it)
222247
Py_ssize_t col_offset = -1;
223248
Py_ssize_t end_col_offset = -1;
224249
if (token.start != NULL && token.start >= line_start) {
225-
col_offset = _PyPegen_byte_offset_to_character_offset(line, token.start - line_start);
250+
Py_ssize_t byte_offset = token.start - line_start;
251+
col_offset = byte_offset - it->byte_col_offset_diff;
226252
}
227253
if (token.end != NULL && token.end >= it->tok->line_start) {
228-
end_col_offset = _PyPegen_byte_offset_to_character_offset_raw(it->tok->line_start, token.end - it->tok->line_start);
254+
Py_ssize_t end_byte_offset = token.end - it->tok->line_start;
255+
if (lineno == end_lineno) {
256+
// If the whole token is at the same line, we can just use the token.start
257+
// buffer for figuring out the new column offset, since using line is not
258+
// performant for very long lines.
259+
Py_ssize_t token_byte_offset = token.end - token.start;
260+
Py_ssize_t token_col_offset = _PyPegen_byte_offset_to_character_offset_raw(
261+
token.start, token_byte_offset
262+
);
263+
end_col_offset = col_offset + token_col_offset;
264+
it->byte_col_offset_diff += token_byte_offset - token_col_offset;
265+
} else {
266+
end_col_offset = _PyPegen_byte_offset_to_character_offset_raw(it->tok->line_start, end_byte_offset);
267+
it->byte_col_offset_diff += end_byte_offset - end_col_offset;
268+
}
229269
}
230270

231271
if (it->tok->tok_extra_tokens) {
@@ -262,7 +302,7 @@ tokenizeriter_next(tokenizeriterobject *it)
262302
}
263303
}
264304

265-
result = Py_BuildValue("(iN(nn)(nn)N)", type, str, lineno, col_offset, end_lineno, end_col_offset, line);
305+
result = Py_BuildValue("(iN(nn)(nn)O)", type, str, lineno, col_offset, end_lineno, end_col_offset, line);
266306
exit:
267307
_PyToken_Free(&token);
268308
if (type == ENDMARKER) {

0 commit comments

Comments
 (0)