Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/asttokens/util.py: 66%
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 collections
17import io
18import re
19import sys
20import token
21import tokenize
22from abc import ABCMeta
23from ast import Module, expr, AST
24from functools import lru_cache
25from typing import (
26 Callable,
27 Dict,
28 Iterable,
29 Iterator,
30 List,
31 Optional,
32 Tuple,
33 Union,
34 cast,
35 Any,
36 TYPE_CHECKING,
37 Type,
38)
40if TYPE_CHECKING: # pragma: no cover
41 from .astroid_compat import NodeNG
42else:
43 NodeNG = Any
46TokenInfo = tokenize.TokenInfo
47_lone_cr_re = re.compile(r'\r(?!\n)')
50def token_repr(tok_type: int, string: Optional[str]) -> str:
51 """Returns a human-friendly representation of a token with the given type and string."""
52 # repr() prefixes unicode with 'u' on Python2 but not Python3; strip it out for consistency.
53 return f"{token.tok_name[tok_type]}:{repr(string).lstrip('u')}"
56class Token(collections.namedtuple('Token', 'type string start end line index startpos endpos')):
57 """
58 TokenInfo is an 8-tuple containing the same 5 fields as the tokens produced by the tokenize
59 module, and 3 additional ones useful for this module:
61 - [0] .type Token type (see token.py)
62 - [1] .string Token (a string)
63 - [2] .start Starting (row, column) indices of the token (a 2-tuple of ints)
64 - [3] .end Ending (row, column) indices of the token (a 2-tuple of ints)
65 - [4] .line Original line (string)
66 - [5] .index Index of the token in the list of tokens that it belongs to.
67 - [6] .startpos Starting character offset into the input text.
68 - [7] .endpos Ending character offset into the input text.
69 """
70 def __str__(self) -> str:
71 return token_repr(self.type, self.string)
74# Type class used to expand out the definition of AST to include fields added by this library
75# It's not actually used for anything other than type checking though!
76class EnhancedAST(AST):
77 # Additional attributes set by mark_tokens
78 first_token: Token = None # type: ignore
79 last_token: Token = None # type: ignore
80 lineno: int = 0
81 end_lineno: int = 0
82 end_col_offset: int = 0
85AstNode = Union[EnhancedAST, NodeNG]
88def match_token(token: Token, tok_type: int, tok_str: Optional[str] = None) -> bool:
89 """Returns true if token is of the given type and, if a string is given, has that string."""
90 return token.type == tok_type and (tok_str is None or token.string == tok_str)
93def expect_token(token: Token, tok_type: int, tok_str: Optional[str] = None) -> None:
94 """
95 Verifies that the given token is of the expected type. If tok_str is given, the token string
96 is verified too. If the token doesn't match, raises an informative ValueError.
97 """
98 if not match_token(token, tok_type, tok_str):
99 raise ValueError(
100 f"Expected token {token_repr(tok_type, tok_str)}, "
101 f"got {str(token)} on line {token.start[0]} col {token.start[1] + 1}"
102 )
105def is_non_coding_token(token_type: int) -> bool:
106 """
107 These are considered non-coding tokens, as they don't affect the syntax tree.
108 """
109 return token_type in (token.NL, token.COMMENT, token.ENCODING)
112def generate_tokens(text: str) -> Iterator[TokenInfo]:
113 """
114 Generates standard library tokens for the given code.
115 """
116 # Before Python 3.12, tokenize treats an entire line containing a lone carriage return as a
117 # non-coding NL token, even though ast.parse treats the carriage return as a line boundary.
118 # Replacing lone carriage returns is length-preserving, so token offsets still map to the
119 # original source. Keep CRLF intact because changing its length would shift later offsets.
120 text = _lone_cr_re.sub('\n', text)
121 # tokenize.generate_tokens is technically an undocumented API for Python3, but allows us to use the same API as for
122 # Python2. See https://stackoverflow.com/a/4952291/328565.
123 # FIXME: Remove cast once https://github.com/python/typeshed/issues/7003 gets fixed
124 return tokenize.generate_tokens(cast(Callable[[], str], io.StringIO(text).readline))
127def iter_children_func(node: AST) -> Callable:
128 """
129 Returns a function which yields all direct children of a AST node,
130 skipping children that are singleton nodes.
131 The function depends on whether ``node`` is from ``ast`` or from the ``astroid`` module.
132 """
133 return iter_children_astroid if hasattr(node, 'get_children') else iter_children_ast
136def iter_children_astroid(node: NodeNG, include_joined_str: bool = False) -> Union[Iterator, list]:
137 if not include_joined_str and is_joined_str(node):
138 return []
140 return node.get_children()
143SINGLETONS = {c for n, c in ast.__dict__.items() if isinstance(c, type) and
144 issubclass(c, (ast.expr_context, ast.boolop, ast.operator, ast.unaryop, ast.cmpop))}
147def iter_children_ast(node: AST, include_joined_str: bool = False) -> Iterator[Union[AST, expr]]:
148 if not include_joined_str and is_joined_str(node):
149 return
151 if isinstance(node, ast.Dict):
152 # override the iteration order: instead of <all keys>, <all values>,
153 # yield keys and values in source order (key1, value1, key2, value2, ...)
154 for (key, value) in zip(node.keys, node.values):
155 if key is not None:
156 yield key
157 yield value
158 return
160 for child in ast.iter_child_nodes(node):
161 # Skip singleton children; they don't reflect particular positions in the code and break the
162 # assumptions about the tree consisting of distinct nodes. Note that collecting classes
163 # beforehand and checking them in a set is faster than using isinstance each time.
164 if child.__class__ not in SINGLETONS:
165 yield child
168stmt_class_names = {n for n, c in ast.__dict__.items()
169 if isinstance(c, type) and issubclass(c, ast.stmt)}
170expr_class_names = ({n for n, c in ast.__dict__.items()
171 if isinstance(c, type) and issubclass(c, ast.expr)} |
172 {'AssignName', 'DelName', 'Const', 'AssignAttr', 'DelAttr'})
174# These feel hacky compared to isinstance() but allow us to work with both ast and astroid nodes
175# in the same way, and without even importing astroid.
176def is_expr(node: AstNode) -> bool:
177 """Returns whether node is an expression node."""
178 return node.__class__.__name__ in expr_class_names
180def is_stmt(node: AstNode) -> bool:
181 """Returns whether node is a statement node."""
182 return node.__class__.__name__ in stmt_class_names
184def is_module(node: AstNode) -> bool:
185 """Returns whether node is a module node."""
186 return node.__class__.__name__ == 'Module'
188def is_joined_str(node: AstNode) -> bool:
189 """Returns whether node is a JoinedStr node, used to represent f-strings."""
190 # At the moment, nodes below JoinedStr have wrong line/col info, and trying to process them only
191 # leads to errors.
192 return node.__class__.__name__ == 'JoinedStr'
195def is_expr_stmt(node: AstNode) -> bool:
196 """Returns whether node is an `Expr` node, which is a statement that is an expression."""
197 return node.__class__.__name__ == 'Expr'
201CONSTANT_CLASSES: Tuple[Type, ...] = (ast.Constant,)
202try:
203 from astroid.nodes import Const
204 CONSTANT_CLASSES += (Const,)
205except ImportError: # pragma: no cover
206 # astroid is not available
207 pass
209def is_constant(node: AstNode) -> bool:
210 """Returns whether node is a Constant node."""
211 return isinstance(node, CONSTANT_CLASSES)
214def is_ellipsis(node: AstNode) -> bool:
215 """Returns whether node is an Ellipsis node."""
216 return is_constant(node) and node.value is Ellipsis # type: ignore
219def is_starred(node: AstNode) -> bool:
220 """Returns whether node is a starred expression node."""
221 return node.__class__.__name__ == 'Starred'
224def is_slice(node: AstNode) -> bool:
225 """Returns whether node represents a slice, e.g. `1:2` in `x[1:2]`"""
226 # Before 3.9, a tuple containing a slice is an ExtSlice,
227 # but this was removed in https://bugs.python.org/issue34822
228 return (
229 node.__class__.__name__ in ('Slice', 'ExtSlice')
230 or (
231 node.__class__.__name__ == 'Tuple'
232 and any(map(is_slice, cast(ast.Tuple, node).elts))
233 )
234 )
237def is_empty_astroid_slice(node: AstNode) -> bool:
238 return (
239 node.__class__.__name__ == "Slice"
240 and not isinstance(node, ast.AST)
241 and node.lower is node.upper is node.step is None
242 )
245# Sentinel value used by visit_tree().
246_PREVISIT = object()
248def visit_tree(
249 node: Module,
250 previsit: Callable[[AstNode, Optional[Token]], Tuple[Optional[Token], Optional[Token]]],
251 postvisit: Optional[Callable[[AstNode, Optional[Token], Optional[Token]], None]]
252) -> None:
253 """
254 Scans the tree under the node depth-first using an explicit stack. It avoids implicit recursion
255 via the function call stack to avoid hitting 'maximum recursion depth exceeded' error.
257 It calls ``previsit()`` and ``postvisit()`` as follows:
259 * ``previsit(node, par_value)`` - should return ``(par_value, value)``
260 ``par_value`` is as returned from ``previsit()`` of the parent.
262 * ``postvisit(node, par_value, value)`` - should return ``value``
263 ``par_value`` is as returned from ``previsit()`` of the parent, and ``value`` is as
264 returned from ``previsit()`` of this node itself. The return ``value`` is ignored except
265 the one for the root node, which is returned from the overall ``visit_tree()`` call.
267 For the initial node, ``par_value`` is None. ``postvisit`` may be None.
268 """
269 if not postvisit:
270 postvisit = lambda node, pvalue, value: None
272 iter_children = iter_children_func(node)
273 done = set()
274 ret = None
275 stack: List[Tuple[AstNode, Optional[Token], object]] = [(node, None, _PREVISIT)]
276 while stack:
277 current, par_value, value = stack.pop()
278 if value is _PREVISIT:
279 assert current not in done # protect againt infinite loop in case of a bad tree.
280 done.add(current)
282 pvalue, post_value = previsit(current, par_value)
283 stack.append((current, par_value, post_value))
285 # Insert all children in reverse order (so that first child ends up on top of the stack).
286 ins = len(stack)
287 for n in iter_children(current):
288 stack.insert(ins, (n, pvalue, _PREVISIT))
289 else:
290 ret = postvisit(current, par_value, cast(Optional[Token], value))
291 return ret
294def walk(node: AST, include_joined_str: bool = False) -> Iterator[Union[Module, AstNode]]:
295 """
296 Recursively yield all descendant nodes in the tree starting at ``node`` (including ``node``
297 itself), using depth-first pre-order traversal (yieling parents before their children).
299 This is similar to ``ast.walk()``, but with a different order, and it works for both ``ast`` and
300 ``astroid`` trees. Also, as ``iter_children()``, it skips singleton nodes generated by ``ast``.
302 By default, ``JoinedStr`` (f-string) nodes and their contents are skipped
303 because they previously couldn't be handled. Set ``include_joined_str`` to True to include them.
304 """
305 iter_children = iter_children_func(node)
306 done = set()
307 stack = [node]
308 while stack:
309 current = stack.pop()
310 assert current not in done # protect againt infinite loop in case of a bad tree.
311 done.add(current)
313 yield current
315 # Insert all children in reverse order (so that first child ends up on top of the stack).
316 # This is faster than building a list and reversing it.
317 ins = len(stack)
318 for c in iter_children(current, include_joined_str):
319 stack.insert(ins, c)
322def replace(text: str, replacements: List[Tuple[int, int, str]]) -> str:
323 """
324 Replaces multiple slices of text with new values. This is a convenience method for making code
325 modifications of ranges e.g. as identified by ``ASTTokens.get_text_range(node)``. Replacements is
326 an iterable of ``(start, end, new_text)`` tuples.
328 For example, ``replace("this is a test", [(0, 4, "X"), (8, 9, "THE")])`` produces
329 ``"X is THE test"``.
330 """
331 p = 0
332 parts = []
333 for (start, end, new_text) in sorted(replacements):
334 parts.append(text[p:start])
335 parts.append(new_text)
336 p = end
337 parts.append(text[p:])
338 return ''.join(parts)
341class NodeMethods:
342 """
343 Helper to get `visit_{node_type}` methods given a node's class and cache the results.
344 """
345 def __init__(self) -> None:
346 self._cache: Dict[Union[ABCMeta, type], Callable[[AstNode, Token, Token], Tuple[Token, Token]]] = {}
348 def get(self, obj: Any, cls: Union[ABCMeta, type]) -> Callable:
349 """
350 Using the lowercase name of the class as node_type, returns `obj.visit_{node_type}`,
351 or `obj.visit_default` if the type-specific method is not found.
352 """
353 method = self._cache.get(cls)
354 if not method:
355 name = "visit_" + cls.__name__.lower()
356 method = getattr(obj, name, obj.visit_default)
357 self._cache[cls] = method
358 return method
361def patched_generate_tokens(original_tokens: Iterable[TokenInfo]) -> Iterator[TokenInfo]:
362 """
363 Fixes tokens yielded by `tokenize.generate_tokens` to handle more non-ASCII characters in identifiers.
364 Workaround for https://github.com/python/cpython/issues/68382.
365 Should only be used when tokenizing a string that is known to be valid syntax,
366 because it assumes that error tokens are not actually errors.
367 Combines groups of consecutive NAME, NUMBER, and/or ERRORTOKEN tokens into a single NAME token.
368 """
369 group: List[tokenize.TokenInfo] = []
370 for tok in original_tokens:
371 if (
372 tok.type in (tokenize.NAME, tokenize.ERRORTOKEN, tokenize.NUMBER)
373 # Only combine tokens if they have no whitespace in between
374 and (not group or group[-1].end == tok.start)
375 ):
376 group.append(tok)
377 else:
378 for combined_token in combine_tokens(group):
379 yield combined_token
380 group = []
381 yield tok
382 for combined_token in combine_tokens(group):
383 yield combined_token
385def combine_tokens(group: List[tokenize.TokenInfo]) -> List[tokenize.TokenInfo]:
386 if not any(tok.type == tokenize.ERRORTOKEN for tok in group) or len({tok.line for tok in group}) != 1:
387 return group
388 return [
389 tokenize.TokenInfo(
390 type=tokenize.NAME,
391 string="".join(t.string for t in group),
392 start=group[0].start,
393 end=group[-1].end,
394 line=group[0].line,
395 )
396 ]
399def last_stmt(node: AstNode) -> AstNode:
400 """
401 If the given AST node contains multiple statements, return the last one.
402 Otherwise, just return the node.
403 """
404 child_stmts = [
405 child for child in iter_children_func(node)(node)
406 if is_stmt(child) or type(child).__name__ in (
407 "excepthandler",
408 "ExceptHandler",
409 "match_case",
410 "MatchCase",
411 "TryExcept",
412 "TryFinally",
413 )
414 ]
415 if child_stmts:
416 return last_stmt(child_stmts[-1])
417 return node
421@lru_cache(maxsize=None)
422def fstring_positions_work() -> bool:
423 """
424 The positions attached to nodes inside f-string FormattedValues have some bugs
425 that were fixed in Python 3.9.7 in https://github.com/python/cpython/pull/27729.
426 This checks for those bugs more concretely without relying on the Python version.
427 Specifically this checks:
428 - Values with a format spec or conversion
429 - Repeated (i.e. identical-looking) expressions
430 - f-strings implicitly concatenated over multiple lines.
431 - Multiline, triple-quoted f-strings.
432 """
433 source = """(
434 f"a {b}{b} c {d!r} e {f:g} h {i:{j}} k {l:{m:n}}"
435 f"a {b}{b} c {d!r} e {f:g} h {i:{j}} k {l:{m:n}}"
436 f"{x + y + z} {x} {y} {z} {z} {z!a} {z:z}"
437 f'''
438 {s} {t}
439 {u} {v}
440 '''
441 )"""
442 tree = ast.parse(source)
443 name_nodes = [node for node in ast.walk(tree) if isinstance(node, ast.Name)]
444 name_positions = [(node.lineno, node.col_offset) for node in name_nodes]
445 positions_are_unique = len(set(name_positions)) == len(name_positions)
446 correct_source_segments = all(
447 ast.get_source_segment(source, node) == node.id
448 for node in name_nodes
449 )
450 return positions_are_unique and correct_source_segments
452def annotate_fstring_nodes(tree: ast.AST) -> None:
453 """
454 Add a special attribute `_broken_positions` to nodes inside f-strings
455 if the lineno/col_offset cannot be trusted.
456 """
457 if sys.version_info >= (3, 12):
458 # f-strings were weirdly implemented until https://peps.python.org/pep-0701/
459 # In Python 3.12, inner nodes have sensible positions.
460 return
461 for joinedstr in walk(tree, include_joined_str=True):
462 if not isinstance(joinedstr, ast.JoinedStr):
463 continue
464 for part in joinedstr.values:
465 # The ast positions of the FormattedValues/Constant nodes span the full f-string, which is weird.
466 setattr(part, '_broken_positions', True) # use setattr for mypy
468 if isinstance(part, ast.FormattedValue):
469 if not fstring_positions_work():
470 for child in walk(part.value):
471 setattr(child, '_broken_positions', True)
473 if part.format_spec: # this is another JoinedStr
474 # Again, the standard positions span the full f-string.
475 setattr(part.format_spec, '_broken_positions', True)