1"""Token-related utilities"""
2
3# Copyright (c) IPython Development Team.
4# Distributed under the terms of the Modified BSD License.
5from __future__ import annotations
6
7import itertools
8import tokenize
9from io import StringIO
10from keyword import iskeyword
11from tokenize import TokenInfo
12from typing import NamedTuple
13from collections.abc import Callable
14from collections.abc import Generator
15
16
17class Token(NamedTuple):
18 token: int
19 text: str
20 start: int
21 end: int
22 line: str
23
24
25def generate_tokens(readline: Callable) -> Generator[TokenInfo]:
26 """wrap generate_tkens to catch EOF errors"""
27 try:
28 yield from tokenize.generate_tokens(readline)
29 except tokenize.TokenError:
30 # catch EOF error
31 return
32
33
34def generate_tokens_catch_errors(
35 readline, extra_errors_to_catch: list[str] | None = None
36):
37 default_errors_to_catch = [
38 "unterminated string literal",
39 "invalid non-printable character",
40 "after line continuation character",
41 ]
42 assert extra_errors_to_catch is None or isinstance(extra_errors_to_catch, list)
43 errors_to_catch = default_errors_to_catch + (extra_errors_to_catch or [])
44
45 tokens: list[TokenInfo] = []
46 try:
47 for token in tokenize.generate_tokens(readline):
48 tokens.append(token)
49 yield token
50 except tokenize.TokenError as exc:
51 if any(error in exc.args[0] for error in errors_to_catch):
52 if tokens:
53 start = tokens[-1].start[0], tokens[-1].end[0]
54 end = start
55 line = tokens[-1].line
56 else:
57 start = end = (1, 0)
58 line = ""
59 yield TokenInfo(tokenize.ERRORTOKEN, "", start, end, line)
60 else:
61 # Catch EOF
62 raise
63
64
65def line_at_cursor(cell: str, cursor_pos: int = 0) -> tuple[str, int]:
66 """Return the line in a cell at a given cursor position
67
68 Used for calling line-based APIs that don't support multi-line input, yet.
69
70 Parameters
71 ----------
72 cell : str
73 multiline block of text
74 cursor_pos : integer
75 the cursor position
76
77 Returns
78 -------
79 (line, offset): (string, integer)
80 The line with the current cursor, and the character offset of the start of the line.
81 """
82 offset = 0
83 lines = cell.splitlines(True)
84 for line in lines:
85 next_offset = offset + len(line)
86 if not line.endswith("\n"):
87 # If the last line doesn't have a trailing newline, treat it as if
88 # it does so that the cursor at the end of the line still counts
89 # as being on that line.
90 next_offset += 1
91 if next_offset > cursor_pos:
92 break
93 offset = next_offset
94 else:
95 line = ""
96 return line, offset
97
98
99def token_at_cursor(cell: str, cursor_pos: int = 0) -> str:
100 """Get the token at a given cursor
101
102 Used for introspection.
103
104 Function calls are prioritized, so the token for the callable will be returned
105 if the cursor is anywhere inside the call.
106
107 Parameters
108 ----------
109 cell : str
110 A block of Python code
111 cursor_pos : int
112 The location of the cursor in the block where the token should be found
113 """
114 names: list[str] = []
115 call_names: list[str] = []
116 closing_call_name: str | None = None
117 most_recent_outer_name: str | None = None
118
119 offsets = {1: 0} # lines start at 1
120 intersects_with_cursor = False
121 cur_token_is_name = False
122 tokens: list[Token | None] = [
123 Token(*tup) for tup in generate_tokens(StringIO(cell).readline)
124 ]
125 if not tokens:
126 return ""
127 for prev_tok, (tok, next_tok) in zip(
128 [None] + tokens, itertools.pairwise(tokens + [None])
129 ):
130 # token, text, start, end, line = tup
131 start_line, start_col = tok.start
132 end_line, end_col = tok.end
133 if end_line + 1 not in offsets:
134 # keep track of offsets for each line
135 lines = tok.line.splitlines(True)
136 for lineno, line in enumerate(lines, start_line + 1):
137 if lineno not in offsets:
138 offsets[lineno] = offsets[lineno - 1] + len(line)
139
140 closing_call_name = None
141
142 offset = offsets[start_line]
143 if offset + start_col > cursor_pos:
144 # current token starts after the cursor,
145 # don't consume it
146 break
147
148 if cur_token_is_name := tok.token == tokenize.NAME and not iskeyword(tok.text):
149 if (
150 names
151 and prev_tok
152 and prev_tok.token == tokenize.OP
153 and prev_tok.text == "."
154 ):
155 names[-1] = "{}.{}".format(names[-1], tok.text)
156 else:
157 names.append(tok.text)
158 if (
159 next_tok is not None
160 and next_tok.token == tokenize.OP
161 and next_tok.text == "="
162 ):
163 # don't inspect the lhs of an assignment
164 names.pop(-1)
165 cur_token_is_name = False
166 if not call_names:
167 most_recent_outer_name = names[-1] if names else None
168 elif tok.token == tokenize.OP:
169 if tok.text == "(" and names:
170 # if we are inside a function call, inspect the function
171 call_names.append(names[-1])
172 elif tok.text == ")" and call_names:
173 # keep track of the most recently popped call_name from the stack
174 closing_call_name = call_names.pop(-1)
175
176 if offsets[end_line] + end_col > cursor_pos:
177 # we found the cursor, stop reading
178 # if the current token intersects directly, use it instead of the call token
179 intersects_with_cursor = offsets[start_line] + start_col <= cursor_pos
180 break
181
182 if cur_token_is_name and intersects_with_cursor:
183 return names[-1]
184 # if the cursor isn't directly over a name token, use the most recent
185 # call name if we can find one
186 elif closing_call_name:
187 # if we're on a ")", use the most recently popped call name
188 return closing_call_name
189 elif call_names:
190 # otherwise, look for the most recent call name in the stack
191 return call_names[-1]
192 elif most_recent_outer_name:
193 # if we've popped all the call names, use the most recently-seen
194 # outer name
195 return most_recent_outer_name
196 elif names:
197 # failing that, use the most recently seen name
198 return names[-1]
199 else:
200 # give up
201 return ""