@@ -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 );
266306exit :
267307 _PyToken_Free (& token );
268308 if (type == ENDMARKER ) {
0 commit comments