Skip to content

Commit bb197e7

Browse files
pablogsalpablogsal
committed
gh-153569: track f-string comments as source spans
1 parent 7432a67 commit bb197e7

7 files changed

Lines changed: 111 additions & 48 deletions

File tree

‎Lib/test/test_fstring.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1686,6 +1686,10 @@ def __repr__(self):
16861686
self.assertEqual(f'{"""a""""#" # outside
16871687
=}', '"""a""""#" \n=\'a#\'')
16881688

1689+
d = {'a#b': 42}
1690+
self.assertEqual(f'''{f"{d["a#b"]}"=}''',
1691+
'f"{d["a#b"]}"=\'42\'')
1692+
16891693
self.assertEqual(f'{ # some comment goes here
16901694
"""hello"""=}', ' \n """hello"""=\'hello\'')
16911695
self.assertEqual(f'{"""# this is not a comment

‎Lib/test/test_tstring.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,5 +295,9 @@ def test_triple_quoted(self):
295295
}'
296296
self.assertEqual(t.interpolations[0].expression, '"""a""""#"')
297297

298+
d = {'a#b': 42}
299+
t = t'''{f"{d["a#b"]}"}'''
300+
self.assertEqual(t.interpolations[0].expression, 'f"{d["a#b"]}"')
301+
298302
if __name__ == '__main__':
299303
unittest.main()

‎Parser/lexer/lexer.c‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,18 @@ _PyLexer_get_normal_mode(struct tok_state *tok, tokenizer_mode* current_tok, str
332332
c = tok_nextc(tok);
333333
}
334334

335+
if (INSIDE_FSTRING(tok) && INSIDE_FSTRING_EXPR(current_tok)) {
336+
const char *comment_end = tok->cur;
337+
if (c == '\n') {
338+
comment_end--;
339+
}
340+
if (_PyLexer_record_ftstring_comment(
341+
tok, tok->start, comment_end) < 0) {
342+
tok->done = E_NOMEM;
343+
return MAKE_TOKEN(ERRORTOKEN);
344+
}
345+
}
346+
335347
if (tok->tok_extra_tokens) {
336348
p = tok->start;
337349
}

‎Parser/lexer/lexer_internal.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,8 @@ TOK_NEXT_MODE(struct tok_state *tok)
4747
int _PyLexer_nextc(struct tok_state *);
4848
void _PyLexer_backup(struct tok_state *, int);
4949
void _PyLexer_update_ftstring_expr(struct tok_state *, char);
50+
int _PyLexer_record_ftstring_comment(
51+
struct tok_state *, const char *, const char *);
5052
int _PyLexer_set_ftstring_expr(struct tok_state *, struct token *, char);
5153
int _PyLexer_check_string_prefixes(struct tok_state *, int, int, int, int, int);
5254
int _PyLexer_scan_number(struct tok_state *, struct token *, int, int);

‎Parser/lexer/state.c‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,12 @@ _PyTokenizer_Free(struct tok_state *tok)
7171
Py_XDECREF(tok->module);
7272
_PyTok_ReaderFree(tok);
7373
_PyTok_SourceClear(&tok->source);
74+
tokenizer_comments *comments = tok->ftstring_comments;
75+
while (comments != NULL) {
76+
tokenizer_comments *previous = comments->previous;
77+
PyMem_Free(comments);
78+
comments = previous;
79+
}
7480
PyMem_Free(tok);
7581
}
7682

‎Parser/lexer/state.h‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,14 @@ typedef struct _tokenizer_mode {
6161
enum string_kind_t string_kind;
6262
} tokenizer_mode;
6363

64+
typedef struct _tokenizer_comments {
65+
struct _tokenizer_comments *previous;
66+
Py_ssize_t count;
67+
Py_ssize_t capacity;
68+
int mode;
69+
_PyTok_Span spans[];
70+
} tokenizer_comments;
71+
6472
/* Tokenizer state */
6573
struct tok_state {
6674
/* Input state; buf <= cur <= inp */
@@ -116,6 +124,7 @@ struct tok_state {
116124
// TODO: Factor this into its own thing
117125
tokenizer_mode tok_mode_stack[MAXFSTRINGLEVEL];
118126
int tok_mode_stack_index;
127+
tokenizer_comments *ftstring_comments;
119128
int tok_extra_tokens;
120129
int comment_newline;
121130
int implicit_newline;

‎Parser/lexer/string.c‎

Lines changed: 74 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -25,53 +25,53 @@ offset_pointer(const struct tok_state *tok, _PyTok_Off offset)
2525
return tok->source.bytes + offset;
2626
}
2727

28-
static Py_ssize_t
29-
strip_expr_comments(const char *expr, Py_ssize_t len, char *result)
28+
static tokenizer_comments *
29+
current_comments(const struct tok_state *tok)
3030
{
31-
Py_ssize_t output = 0;
32-
char quote = 0;
33-
int quote_size = 0;
34-
35-
for (Py_ssize_t i = 0; i < len;) {
36-
char c = expr[i];
37-
if (quote != 0) {
38-
if (c == '\\' && i + 1 < len) {
39-
result[output] = c;
40-
result[output + 1] = expr[i + 1];
41-
output += 2;
42-
i += 2;
43-
continue;
44-
}
45-
if (c == quote) {
46-
if (quote_size == 1) {
47-
quote = 0;
48-
}
49-
else if (i + 2 < len && expr[i + 1] == quote &&
50-
expr[i + 2] == quote) {
51-
memcpy(result + output, expr + i, 3);
52-
output += 3;
53-
i += 3;
54-
quote = 0;
55-
continue;
56-
}
57-
}
31+
tokenizer_comments *comments = tok->ftstring_comments;
32+
return comments != NULL && comments->mode == tok->tok_mode_stack_index
33+
? comments : NULL;
34+
}
35+
36+
int
37+
_PyLexer_record_ftstring_comment(struct tok_state *tok, const char *start,
38+
const char *end)
39+
{
40+
tokenizer_mode *mode = TOK_GET_MODE(tok);
41+
if (mode->debug_expr.end >= 0) {
42+
return 0;
43+
}
44+
assert(mode->debug_expr.start >= 0);
45+
tokenizer_comments *comments = current_comments(tok);
46+
if (comments == NULL || comments->count == comments->capacity) {
47+
int create = comments == NULL;
48+
Py_ssize_t max_capacity = (PY_SSIZE_T_MAX -
49+
(Py_ssize_t)sizeof(*comments)) /
50+
(Py_ssize_t)sizeof(*comments->spans);
51+
if (comments != NULL && comments->capacity > max_capacity / 2) {
52+
PyErr_NoMemory();
53+
return -1;
5854
}
59-
else if (c == '#') {
60-
while (i < len && expr[i] != '\n') {
61-
i++;
62-
}
63-
continue;
55+
Py_ssize_t capacity = comments == NULL ? 4 : comments->capacity * 2;
56+
size_t size = sizeof(*comments) +
57+
(size_t)capacity * sizeof(*comments->spans);
58+
tokenizer_comments *resized = PyMem_Realloc(comments, size);
59+
if (resized == NULL) {
60+
PyErr_NoMemory();
61+
return -1;
6462
}
65-
else if (c == '\'' || c == '"') {
66-
quote = c;
67-
quote_size = i + 2 < len && expr[i + 1] == c &&
68-
expr[i + 2] == c ? 3 : 1;
63+
comments = resized;
64+
if (create) {
65+
comments->previous = tok->ftstring_comments;
66+
comments->count = 0;
67+
comments->mode = tok->tok_mode_stack_index;
6968
}
70-
result[output] = c;
71-
output++;
72-
i++;
69+
comments->capacity = capacity;
70+
tok->ftstring_comments = comments;
7371
}
74-
return output;
72+
comments->spans[comments->count++] = (_PyTok_Span){
73+
current_offset(tok, start), current_offset(tok, end)};
74+
return 0;
7575
}
7676

7777
int
@@ -89,21 +89,38 @@ _PyLexer_set_ftstring_expr(struct tok_state* tok, struct token *token, char c) {
8989
if (expr == NULL) {
9090
return -1;
9191
}
92+
tokenizer_comments *comments = current_comments(tok);
9293
PyObject *res;
93-
if (memchr(expr, '#', expr_len) == NULL) {
94-
res = PyUnicode_DecodeUTF8(expr, expr_len, NULL);
95-
}
96-
else {
94+
if (comments != NULL && comments->count > 0) {
9795
char *stripped = PyMem_Malloc((size_t)expr_len);
9896
if (stripped == NULL) {
9997
PyErr_NoMemory();
10098
return -1;
10199
}
102-
Py_ssize_t stripped_len = strip_expr_comments(
103-
expr, expr_len, stripped);
100+
_PyTok_Off copied_to = tok_mode->debug_expr.start;
101+
Py_ssize_t stripped_len = 0;
102+
for (Py_ssize_t i = 0; i < comments->count; i++) {
103+
_PyTok_Span comment = comments->spans[i];
104+
assert(comment.start >= copied_to);
105+
assert(comment.end <= tok_mode->debug_expr.end);
106+
Py_ssize_t length = comment.start - copied_to;
107+
memcpy(stripped + stripped_len,
108+
expr + copied_to - tok_mode->debug_expr.start,
109+
(size_t)length);
110+
stripped_len += length;
111+
copied_to = comment.end;
112+
}
113+
Py_ssize_t length = tok_mode->debug_expr.end - copied_to;
114+
memcpy(stripped + stripped_len,
115+
expr + copied_to - tok_mode->debug_expr.start,
116+
(size_t)length);
117+
stripped_len += length;
104118
res = PyUnicode_DecodeUTF8(stripped, stripped_len, NULL);
105119
PyMem_Free(stripped);
106120
}
121+
else {
122+
res = PyUnicode_DecodeUTF8(expr, expr_len, NULL);
123+
}
107124

108125
if (!res) {
109126
return -1;
@@ -121,6 +138,10 @@ _PyLexer_update_ftstring_expr(struct tok_state *tok, char cur)
121138
case '{':
122139
tok_mode->debug_expr = (_PyTok_Span){
123140
current_offset(tok, tok->cur), -1};
141+
tokenizer_comments *comments = current_comments(tok);
142+
if (comments != NULL) {
143+
comments->count = 0;
144+
}
124145
break;
125146
case '}':
126147
case '!':
@@ -419,6 +440,11 @@ _PyLexer_get_fstring_mode(struct tok_state *tok, tokenizer_mode* current_tok, st
419440

420441
p_start = tok->start;
421442
p_end = tok->cur;
443+
tokenizer_comments *comments = current_comments(tok);
444+
if (comments != NULL) {
445+
tok->ftstring_comments = comments->previous;
446+
PyMem_Free(comments);
447+
}
422448
tok->tok_mode_stack_index--;
423449
return MAKE_TOKEN(FTSTRING_END(current_tok));
424450

0 commit comments

Comments
 (0)