gh-153569: include token types in tokenizer results

Author: pablogsal

Parser/lexer/lexer.c

@@ -511,9 +511,10 @@ _PyLexer_get_normal(struct tok_state *tok, ftstring_state *current, struct token
 }
 
 
-int
+void
 _PyTokenizer_Get(struct tok_state *tok, struct token *token)
 {
+    _PyToken_Free(token);
     ftstring_state *current = _PyLexer_CurrentFTString(tok);
     int result;
     if (current == NULL) {
@@ -535,5 +536,5 @@ _PyTokenizer_Get(struct tok_state *tok, struct token *token)
     if (tok_failed(tok)) {
         result = ERRORTOKEN;
     }
-    return result;
+    token->type = result;
 }

Parser/lexer/state.c

@@ -74,19 +74,14 @@ _PyTokenizer_Free(struct tok_state *tok)
     PyMem_Free(tok);
 }
 
-void
-_PyToken_Free(struct token *token) {
-    Py_XDECREF(token->metadata);
-}
-
 void
 _PyToken_Init(struct token *token) {
-#ifdef Py_DEBUG
-    token->span = (_PyTok_Span){-1, -1};
-    token->start_loc = (_PyTok_Loc){-1, -1};
-    token->end_loc = (_PyTok_Loc){-1, -1};
-#endif
-    token->metadata = NULL;
+    *token = (struct token){
+        .type = -1,
+        .span = {-1, -1},
+        .start_loc = {-1, -1},
+        .end_loc = {-1, -1},
+    };
 }
 
 int

Parser/pegen.c

@@ -206,7 +206,6 @@ initialize_token(Parser *p, Token *parser_token, struct token *new_token, int to
     parser_token->metadata = NULL;
     if (new_token->metadata != NULL) {
         if (_PyArena_AddPyObject(p->arena, new_token->metadata) < 0) {
-            Py_DECREF(new_token->metadata);
             return -1;
         }
         parser_token->metadata = new_token->metadata;
@@ -260,10 +259,10 @@ _PyPegen_fill_token(Parser *p)
 {
     struct token new_token;
     _PyToken_Init(&new_token);
-    int type = _PyTokenizer_Get(p->tok, &new_token);
+    _PyTokenizer_Get(p->tok, &new_token);
 
     // Record and skip '# type: ignore' comments
-    while (type == TYPE_IGNORE) {
+    while (new_token.type == TYPE_IGNORE) {
         Py_ssize_t len;
         const char *text = _PyToken_TextView(p->tok, &new_token, &len);
         char *tag = PyMem_Malloc((size_t)len + 1);
@@ -278,9 +277,11 @@ _PyPegen_fill_token(Parser *p)
             PyErr_NoMemory();
             goto error;
         }
-        type = _PyTokenizer_Get(p->tok, &new_token);
+        _PyTokenizer_Get(p->tok, &new_token);
     }
 
+    int type = new_token.type;
+
     // If we have reached the end and we are in single input mode we need to insert a newline and reset the parsing
     if (p->start_rule == Py_single_input && type == ENDMARKER && p->parsing_started) {
         type = NEWLINE; /* Add an extra newline */
@@ -300,7 +301,9 @@ _PyPegen_fill_token(Parser *p)
     }
 
     Token *t = p->tokens[p->fill];
-    return initialize_token(p, t, &new_token, type);
+    int result = initialize_token(p, t, &new_token, type);
+    _PyToken_Free(&new_token);
+    return result;
 error:
     _PyToken_Free(&new_token);
     return -1;

Parser/pegen_errors.c

@@ -138,7 +138,8 @@ _PyPegen_tokenize_full_source_to_check_for_errors(Parser *p) {
     _PyToken_Init(&new_token);
 
     for (;;) {
-        switch (_PyTokenizer_Get(p->tok, &new_token)) {
+        _PyTokenizer_Get(p->tok, &new_token);
+        switch (new_token.type) {
             case ERRORTOKEN: {
                 if (PyErr_Occurred()) {
                     ret = -1;

Parser/tokenizer/api.c

@@ -63,7 +63,7 @@ _PyTokenizer_SpanView(const struct tok_state *tok, _PyTok_Span span,
 
 void
 _PyToken_GetView(const struct tok_state *tok, const struct token *token,
-                 int type, _PyToken_View *view)
+                 _PyToken_View *view)
 {
     assert(view != NULL);
     assert((token->span.start == -1 && token->span.end == -1) ||
@@ -75,7 +75,7 @@ _PyToken_GetView(const struct tok_state *tok, const struct token *token,
         ? NULL : _PyLexer_BufferPointer(tok, token->span.start);
     view->length = token->span.end - token->span.start;
     view->end_line = _PyLexer_BufferPointer(tok, tok->line_start);
-    view->line = ISSTRINGLIT(type)
+    view->line = ISSTRINGLIT(token->type)
         ? view->text - token->start_loc.byte_col : view->end_line;
     view->line_length = tok->inp - tok->line_start +
         (view->end_line - view->line);

Parser/tokenizer/tokenizer.h

@@ -9,6 +9,7 @@ struct tok_state;
 /* Initialize before use. metadata owns a reference released by _PyToken_Free;
    a consumer taking that reference must set metadata to NULL. */
 struct token {
+    int type;
     int level;
     int is_raw;
     _PyTok_Span span;
@@ -52,11 +53,17 @@ typedef struct {
     const char *encoding;
 } _PyTokenizer_Info;
 
-int _PyTokenizer_Get(struct tok_state *, struct token *);
+/* Get replaces the initialized token, releasing its previous metadata.
+   Errors are returned as ERRORTOKEN, with or without a Python exception. */
+void _PyTokenizer_Get(struct tok_state *, struct token *);
 void _PyTokenizer_Free(struct tok_state *);
 void _PyTokenizer_raise_init_error(PyObject *filename);
 void _PyToken_Init(struct token *);
-void _PyToken_Free(struct token *);
+static inline void
+_PyToken_Free(struct token *token)
+{
+    Py_CLEAR(token->metadata);
+}
 
 /* Views and borrowed snapshot references remain valid until the tokenizer is
    mutated or freed. Source spans may be discarded when reading more input. */
@@ -65,10 +72,10 @@ _PyTokenizer_Diagnostic _PyTokenizer_GetDiagnostic(const struct tok_state *);
 /* An absent token span has a nonnull empty text view. */
 const char *_PyToken_TextView(
     const struct tok_state *, const struct token *, Py_ssize_t *);
-/* Pair the token with the type returned by the most recent Get. text is NULL
-   for an absent span; line includes the token's complete physical line range. */
+/* Use the token from the most recent Get. text is NULL for an absent span;
+   line includes the token's complete physical line range. */
 void _PyToken_GetView(
-    const struct tok_state *tok, const struct token *token, int type,
+    const struct tok_state *tok, const struct token *token,
     _PyToken_View *view);
 const char *_PyTokenizer_SpanView(
     const struct tok_state *, _PyTok_Span, Py_ssize_t *);

Python/Python-tokenize.c

@@ -257,7 +257,8 @@ tokenizeriter_next(PyObject *op)
     struct token token;
     _PyToken_Init(&token);
 
-    int type = _PyTokenizer_Get(it->tok, &token);
+    _PyTokenizer_Get(it->tok, &token);
+    int type = token.type;
     if (type == ERRORTOKEN) {
         if(!PyErr_Occurred()) {
             _tokenizer_error(it);
@@ -271,7 +272,7 @@ tokenizeriter_next(PyObject *op)
         goto exit;
     }
     _PyToken_View view;
-    _PyToken_GetView(it->tok, &token, type, &view);
+    _PyToken_GetView(it->tok, &token, &view);
     const char *token_start = view.text;
     PyObject *str;
     if (token.span.start < 0) {