Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright 2016 Grist Labs, Inc.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# https://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15import ast
16import numbers
17import sys
18import token
19from ast import Module
20from typing import Callable, List, Union, cast, Optional, Tuple, TYPE_CHECKING
22from . import util
23from .asttokens import ASTTokens
24from .astroid_compat import astroid_node_classes as nc, BaseContainer as AstroidBaseContainer
25from .util import AstNode
28# Mapping of matching braces. To find a token here, look up token[:2].
29_matching_pairs_left = {
30 (token.OP, '('): (token.OP, ')'),
31 (token.OP, '['): (token.OP, ']'),
32 (token.OP, '{'): (token.OP, '}'),
33}
35_matching_pairs_right = {
36 (token.OP, ')'): (token.OP, '('),
37 (token.OP, ']'): (token.OP, '['),
38 (token.OP, '}'): (token.OP, '{'),
39}
42class MarkTokens:
43 """
44 Helper that visits all nodes in the AST tree and assigns .first_token and .last_token attributes
45 to each of them. This is the heart of the token-marking logic.
46 """
47 def __init__(self, code: ASTTokens) -> None:
48 self._code = code
49 self._methods = util.NodeMethods()
50 self._iter_children: Optional[Callable] = None
52 def visit_tree(self, node: Module) -> None:
53 self._iter_children = util.iter_children_func(node)
54 util.visit_tree(node, self._visit_before_children, self._visit_after_children)
56 def _visit_before_children(
57 self, node: AstNode, parent_token: Optional[util.Token]
58 ) -> Tuple[Optional[util.Token], Optional[util.Token]]:
59 col = getattr(node, 'col_offset', None)
60 token = self._code.get_token_from_utf8(node.lineno, col) if col is not None else None
62 if not token and util.is_module(node):
63 # We'll assume that a Module node starts at the start of the source code.
64 token = self._code.get_token(1, 0)
66 # Use our own token, or our parent's if we don't have one, to pass to child calls as
67 # parent_token argument. The second value becomes the token argument of _visit_after_children.
68 return (token or parent_token, token)
70 def _visit_after_children(
71 self, node: AstNode, parent_token: Optional[util.Token], token: Optional[util.Token]
72 ) -> None:
73 # This processes the node generically first, after all children have been processed.
75 # Get the first and last tokens that belong to children. Note how this doesn't assume that we
76 # iterate through children in order that corresponds to occurrence in source code. This
77 # assumption can fail (e.g. with return annotations).
78 first = token
79 last = None
80 for child in cast(Callable, self._iter_children)(node):
81 # astroid slices have especially wrong positions, we don't want them to corrupt their parents.
82 if util.is_empty_astroid_slice(child):
83 continue
84 if not first or child.first_token.index < first.index:
85 first = child.first_token
86 if not last or child.last_token.index > last.index:
87 last = child.last_token
89 # If we don't have a first token from _visit_before_children, and there were no children, then
90 # use the parent's token as the first token.
91 first = first or parent_token
93 # If no children, set last token to the first one.
94 last = last or first
96 # Statements continue to before NEWLINE. This helps cover a few different cases at once.
97 if util.is_stmt(node):
98 last = self._find_last_in_stmt(cast(util.Token, last))
100 # Capture any unmatched brackets.
101 first, last = self._expand_to_matching_pairs(cast(util.Token, first), cast(util.Token, last), node)
103 # Give a chance to node-specific methods to adjust.
104 nfirst, nlast = self._methods.get(self, node.__class__)(node, first, last)
106 if (nfirst, nlast) != (first, last):
107 # If anything changed, expand again to capture any unmatched brackets.
108 nfirst, nlast = self._expand_to_matching_pairs(nfirst, nlast, node)
110 node.first_token = nfirst
111 node.last_token = nlast
113 def _find_last_in_stmt(self, start_token: util.Token) -> util.Token:
114 t = start_token
115 while (not util.match_token(t, token.NEWLINE) and
116 not util.match_token(t, token.OP, ';') and
117 not token.ISEOF(t.type)):
118 t = self._code.next_token(t, include_extra=True)
119 return self._code.prev_token(t)
121 def _expand_to_matching_pairs(
122 self, first_token: util.Token, last_token: util.Token, node: AstNode
123 ) -> Tuple[util.Token, util.Token]:
124 """
125 Scan tokens in [first_token, last_token] range that are between node's children, and for any
126 unmatched brackets, adjust first/last tokens to include the closing pair.
127 """
128 # We look for opening parens/braces among non-child tokens (i.e. tokens between our actual
129 # child nodes). If we find any closing ones, we match them to the opens.
130 to_match_right: List[Tuple[int, str]] = []
131 to_match_left = []
132 for tok in self._code.token_range(first_token, last_token):
133 tok_info = tok[:2]
134 if to_match_right and tok_info == to_match_right[-1]:
135 to_match_right.pop()
136 elif tok_info in _matching_pairs_left:
137 to_match_right.append(_matching_pairs_left[tok_info])
138 elif tok_info in _matching_pairs_right:
139 to_match_left.append(_matching_pairs_right[tok_info])
141 # Once done, extend `last_token` to match any unclosed parens/braces.
142 for match in reversed(to_match_right):
143 last = self._code.next_token(last_token)
144 # Allow for trailing commas or colons (allowed in subscripts) before the closing delimiter
145 while any(util.match_token(last, token.OP, x) for x in (',', ':')):
146 last = self._code.next_token(last)
147 # Now check for the actual closing delimiter.
148 if util.match_token(last, *match):
149 last_token = last
151 # And extend `first_token` to match any unclosed opening parens/braces.
152 for match in to_match_left:
153 first = self._code.prev_token(first_token)
154 if util.match_token(first, *match):
155 first_token = first
157 return (first_token, last_token)
159 #----------------------------------------------------------------------
160 # Node visitors. Each takes a preliminary first and last tokens, and returns the adjusted pair
161 # that will actually be assigned.
163 def visit_default(
164 self, node: AstNode, first_token: util.Token, last_token: util.Token
165 ) -> Tuple[util.Token, util.Token]:
166 # pylint: disable=no-self-use
167 # By default, we don't need to adjust the token we computed earlier.
168 return (first_token, last_token)
170 def handle_comp(
171 self, open_brace: str, node: AstNode, first_token: util.Token, last_token: util.Token
172 ) -> Tuple[util.Token, util.Token]:
173 # For list/set/dict comprehensions, we only get the token of the first child, so adjust it to
174 # include the opening brace (the closing brace will be matched automatically).
175 before = self._code.prev_token(first_token)
176 util.expect_token(before, token.OP, open_brace)
177 return (before, last_token)
179 def visit_comprehension(
180 self,
181 node: AstNode,
182 first_token: util.Token,
183 last_token: util.Token,
184 ) -> Tuple[util.Token, util.Token]:
185 # The 'comprehension' node starts with 'for' but we only get first child; we search backwards
186 # to find the 'for' keyword.
187 first = self._code.find_token(first_token, token.NAME, 'for', reverse=True)
188 return (first, last_token)
190 def visit_if(
191 self, node: util.Token, first_token: util.Token, last_token: util.Token
192 ) -> Tuple[util.Token, util.Token]:
193 while first_token.string not in ('if', 'elif'):
194 first_token = self._code.prev_token(first_token)
195 return first_token, last_token
197 def handle_attr(
198 self, node: AstNode, first_token: util.Token, last_token: util.Token
199 ) -> Tuple[util.Token, util.Token]:
200 # Attribute node has ".attr" (2 tokens) after the last child.
201 dot = self._code.find_token(last_token, token.OP, '.')
202 name = self._code.next_token(dot)
203 util.expect_token(name, token.NAME)
204 return (first_token, name)
206 visit_attribute = handle_attr
207 visit_assignattr = handle_attr
208 visit_delattr = handle_attr
210 def handle_def(
211 self, node: AstNode, first_token: util.Token, last_token: util.Token
212 ) -> Tuple[util.Token, util.Token]:
213 # With astroid, nodes that start with a doc-string can have an empty body, in which case we
214 # need to adjust the last token to include the doc string.
215 if not node.body and (getattr(node, 'doc_node', None) or getattr(node, 'doc', None)): # type: ignore[union-attr]
216 last_token = self._code.find_token(last_token, token.STRING)
218 # Include @ from decorator
219 if first_token.index > 0:
220 prev = self._code.prev_token(first_token)
221 if util.match_token(prev, token.OP, '@'):
222 first_token = prev
223 return (first_token, last_token)
225 visit_classdef = handle_def
226 visit_functiondef = handle_def
228 def handle_following_brackets(
229 self, node: AstNode, last_token: util.Token, opening_bracket: str
230 ) -> util.Token:
231 # This is for calls and subscripts, which have a pair of brackets
232 # at the end which may contain no nodes, e.g. foo() or bar[:].
233 # We look for the opening bracket and then let the matching pair be found automatically
234 # Remember that last_token is at the end of all children,
235 # so we are not worried about encountering a bracket that belongs to a child.
236 first_child = next(cast(Callable, self._iter_children)(node))
237 call_start = self._code.find_token(first_child.last_token, token.OP, opening_bracket)
238 if call_start.index > last_token.index:
239 last_token = call_start
240 return last_token
242 def visit_call(
243 self, node: util.Token, first_token: util.Token, last_token: util.Token
244 ) -> Tuple[util.Token, util.Token]:
245 last_token = self.handle_following_brackets(node, last_token, '(')
247 # Handling a python bug with decorators with empty parens, e.g.
248 # @deco()
249 # def ...
250 if util.match_token(first_token, token.OP, '@'):
251 first_token = self._code.next_token(first_token)
252 return (first_token, last_token)
254 def visit_matchclass(
255 self, node: util.Token, first_token: util.Token, last_token: util.Token
256 ) -> Tuple[util.Token, util.Token]:
257 last_token = self.handle_following_brackets(node, last_token, '(')
258 return (first_token, last_token)
260 def visit_subscript(
261 self,
262 node: AstNode,
263 first_token: util.Token,
264 last_token: util.Token,
265 ) -> Tuple[util.Token, util.Token]:
266 last_token = self.handle_following_brackets(node, last_token, '[')
267 return (first_token, last_token)
269 def visit_slice(
270 self, node: AstNode, first_token: util.Token, last_token: util.Token
271 ) -> Tuple[util.Token, util.Token]:
272 # consume `:` tokens to the left and right. In Python 3.9, Slice nodes are
273 # given a col_offset, (and end_col_offset), so this will always start inside
274 # the slice, even if it is the empty slice. However, in 3.8 and below, this
275 # will only expand to the full slice if the slice contains a node with a
276 # col_offset. So x[:] will only get the correct tokens in 3.9, but x[1:] and
277 # x[:1] will even on earlier versions of Python.
278 while True:
279 prev = self._code.prev_token(first_token)
280 if prev.string != ':':
281 break
282 first_token = prev
283 while True:
284 next_ = self._code.next_token(last_token)
285 if next_.string != ':':
286 break
287 last_token = next_
288 return (first_token, last_token)
290 def handle_bare_tuple(
291 self, node: AstNode, first_token: util.Token, last_token: util.Token
292 ) -> Tuple[util.Token, util.Token]:
293 # A bare tuple doesn't include parens; if there is a trailing comma, make it part of the tuple.
294 maybe_comma = self._code.next_token(last_token)
295 if util.match_token(maybe_comma, token.OP, ','):
296 last_token = maybe_comma
297 return (first_token, last_token)
299 # In Python3.8 parsed tuples include parentheses when present.
300 def handle_tuple_nonempty(
301 self, node: AstNode, first_token: util.Token, last_token: util.Token
302 ) -> Tuple[util.Token, util.Token]:
303 assert isinstance(node, ast.Tuple) or isinstance(node, AstroidBaseContainer)
304 # It's a bare tuple if the first token belongs to the first child. The first child may
305 # include extraneous parentheses (which don't create new nodes), so account for those too.
306 child = node.elts[0]
307 if TYPE_CHECKING:
308 child = cast(AstNode, child)
309 child_first, child_last = self._gobble_parens(child.first_token, child.last_token, True)
310 if first_token == child_first:
311 return self.handle_bare_tuple(node, first_token, last_token)
312 return (first_token, last_token)
314 def visit_tuple(
315 self, node: AstNode, first_token: util.Token, last_token: util.Token
316 ) -> Tuple[util.Token, util.Token]:
317 assert isinstance(node, ast.Tuple) or isinstance(node, AstroidBaseContainer)
318 if not node.elts:
319 # An empty tuple is just "()", and we need no further info.
320 return (first_token, last_token)
321 return self.handle_tuple_nonempty(node, first_token, last_token)
323 def _gobble_parens(
324 self, first_token: util.Token, last_token: util.Token, include_all: bool = False
325 ) -> Tuple[util.Token, util.Token]:
326 # Expands a range of tokens to include one or all pairs of surrounding parentheses, and
327 # returns (first, last) tokens that include these parens.
328 while first_token.index > 0:
329 prev = self._code.prev_token(first_token)
330 next = self._code.next_token(last_token)
331 if util.match_token(prev, token.OP, '(') and util.match_token(next, token.OP, ')'):
332 first_token, last_token = prev, next
333 if include_all:
334 continue
335 break
336 return (first_token, last_token)
338 def visit_str(
339 self, node: AstNode, first_token: util.Token, last_token: util.Token
340 ) -> Tuple[util.Token, util.Token]:
341 return self.handle_str(first_token, last_token)
343 def visit_joinedstr(
344 self,
345 node: AstNode,
346 first_token: util.Token,
347 last_token: util.Token,
348 ) -> Tuple[util.Token, util.Token]:
349 if sys.version_info < (3, 12):
350 # Older versions don't tokenize the contents of f-strings
351 return self.handle_str(first_token, last_token)
353 last = first_token
354 while True:
355 if util.match_token(last, getattr(token, "FSTRING_START")):
356 # Python 3.12+ has tokens for the start (e.g. `f"`) and end (`"`)
357 # of the f-string. We can't just look for the next FSTRING_END
358 # because f-strings can be nested, e.g. f"{f'{x}'}", so we need
359 # to treat this like matching balanced parentheses.
360 count = 1
361 while count > 0:
362 last = self._code.next_token(last)
363 # mypy complains about token.FSTRING_START and token.FSTRING_END.
364 if util.match_token(last, getattr(token, "FSTRING_START")):
365 count += 1
366 elif util.match_token(last, getattr(token, "FSTRING_END")):
367 count -= 1
368 last_token = last
369 last = self._code.next_token(last_token)
370 elif util.match_token(last, token.STRING):
371 # Similar to handle_str, we also need to handle adjacent strings.
372 last_token = last
373 last = self._code.next_token(last_token)
374 else:
375 break
376 return (first_token, last_token)
378 def visit_bytes(
379 self, node: AstNode, first_token: util.Token, last_token: util.Token
380 ) -> Tuple[util.Token, util.Token]:
381 return self.handle_str(first_token, last_token)
383 def handle_str(
384 self, first_token: util.Token, last_token: util.Token
385 ) -> Tuple[util.Token, util.Token]:
386 # Multiple adjacent STRING tokens form a single string.
387 last = self._code.next_token(last_token)
388 while util.match_token(last, token.STRING):
389 last_token = last
390 last = self._code.next_token(last_token)
391 return (first_token, last_token)
393 def handle_num(
394 self,
395 node: AstNode,
396 value: Union[complex, int, numbers.Number],
397 first_token: util.Token,
398 last_token: util.Token,
399 ) -> Tuple[util.Token, util.Token]:
400 # A constant like '-1' gets turned into two tokens; this will skip the '-'.
401 while util.match_token(last_token, token.OP):
402 last_token = self._code.next_token(last_token)
404 if isinstance(value, complex):
405 # A complex number like -2j cannot be compared directly to 0
406 # A complex number like 1-2j is expressed as a binary operation
407 # so we don't need to worry about it
408 value = value.imag
410 # This makes sure that the - is included
411 if value < 0 and first_token.type == token.NUMBER: # type: ignore[operator]
412 first_token = self._code.prev_token(first_token)
413 return (first_token, last_token)
415 def visit_num(
416 self, node: AstNode, first_token: util.Token, last_token: util.Token
417 ) -> Tuple[util.Token, util.Token]:
418 n = node.n # type: ignore[union-attr] # ast.Num has been removed in python 3.14
419 assert isinstance(n, (complex, int, numbers.Number))
420 return self.handle_num(node, n, first_token, last_token)
422 def visit_const(
423 self, node: AstNode, first_token: util.Token, last_token: util.Token
424 ) -> Tuple[util.Token, util.Token]:
425 assert isinstance(node, ast.Constant) or isinstance(node, nc.Const)
426 if isinstance(node.value, numbers.Number):
427 return self.handle_num(node, node.value, first_token, last_token)
428 elif isinstance(node.value, (str, bytes)):
429 return self.visit_str(node, first_token, last_token)
430 return (first_token, last_token)
432 visit_constant = visit_const
434 def visit_keyword(
435 self, node: AstNode, first_token: util.Token, last_token: util.Token
436 ) -> Tuple[util.Token, util.Token]:
437 # Until python 3.9 (https://bugs.python.org/issue40141),
438 # ast.keyword nodes didn't have line info. Astroid has lineno None.
439 assert isinstance(node, ast.keyword) or isinstance(node, nc.Keyword)
440 if node.arg is not None and getattr(node, 'lineno', None) is None:
441 equals = self._code.find_token(first_token, token.OP, '=', reverse=True)
442 name = self._code.prev_token(equals)
443 util.expect_token(name, token.NAME, node.arg)
444 first_token = name
445 return (first_token, last_token)
447 def visit_starred(
448 self, node: AstNode, first_token: util.Token, last_token: util.Token
449 ) -> Tuple[util.Token, util.Token]:
450 # Astroid has 'Starred' nodes (for "foo(*bar)" type args), but they need to be adjusted.
451 if not util.match_token(first_token, token.OP, '*'):
452 star = self._code.prev_token(first_token)
453 if util.match_token(star, token.OP, '*'):
454 first_token = star
455 return (first_token, last_token)
457 def visit_assignname(
458 self, node: AstNode, first_token: util.Token, last_token: util.Token
459 ) -> Tuple[util.Token, util.Token]:
460 # Astroid may turn 'except' clause into AssignName, but we need to adjust it.
461 if util.match_token(first_token, token.NAME, 'except'):
462 colon = self._code.find_token(last_token, token.OP, ':')
463 first_token = last_token = self._code.prev_token(colon)
464 return (first_token, last_token)
466 # Async nodes should typically start with the word 'async'
467 # but Python < 3.7 doesn't put the col_offset there
468 # AsyncFunctionDef is slightly different because it might have
469 # decorators before that, which visit_functiondef handles
470 def handle_async(
471 self, node: AstNode, first_token: util.Token, last_token: util.Token
472 ) -> Tuple[util.Token, util.Token]:
473 if not first_token.string == 'async':
474 first_token = self._code.prev_token(first_token)
475 return (first_token, last_token)
477 visit_asyncfor = handle_async
478 visit_asyncwith = handle_async
480 def visit_asyncfunctiondef(
481 self,
482 node: AstNode,
483 first_token: util.Token,
484 last_token: util.Token,
485 ) -> Tuple[util.Token, util.Token]:
486 if util.match_token(first_token, token.NAME, 'def'):
487 # Include the 'async' token
488 first_token = self._code.prev_token(first_token)
489 return self.visit_functiondef(node, first_token, last_token)