diff --git a/executing/_linenos.py b/executing/_linenos.py new file mode 100644 index 0000000..c01749f --- /dev/null +++ b/executing/_linenos.py @@ -0,0 +1,59 @@ +import ast + + +def pos_range(node): + if isinstance(node, ast.Module): + start = pos_range(node.body[0])[0] + end = pos_range(node.body[-1])[1] + return start, end + start, end = node, node + + if hasattr(node, "decorator_list") and node.decorator_list: + start = node.decorator_list[0] + + return (start.lineno, start.col_offset), (end.end_lineno, end.end_col_offset) + + +def childs(node): + for child in ast.iter_child_nodes(node): + if not hasattr(child, "lineno"): + for c in childs(child): + yield c + else: + yield child + + +class LinenosCache: + def __init__(self, tree): + self.tree = tree + self.cache = {} + + def __getitem__(self, line): + if line in self.cache: + return self.cache[line] + + result = [] + + def line_items(node): + start, end = pos_range(node) + + if hasattr(node, "lineno"): + if ( + hasattr(node, "end_lineno") + and isinstance(node, ast.expr) + and node.lineno <= line <= node.end_lineno + ): + result.append(node) + elif node.lineno == line: + if hasattr(node, "lineno"): + result.append(node) + + if start[0] <= line <= end[0]: + for child in childs(node): + line_items(child) + + line_items(self.tree) + + self.cache[line] = result + + return result diff --git a/executing/_position_node_finder.py b/executing/_position_node_finder.py index e5cddc3..f463dc3 100644 --- a/executing/_position_node_finder.py +++ b/executing/_position_node_finder.py @@ -5,6 +5,7 @@ from typing import Any, Callable, Iterator, Optional, Sequence, Set, Tuple, Type, Union, cast from .executing import EnhancedAST, NotOneValueFound, Source, only, function_node_types, assert_ from ._exceptions import KnownIssue, VerifierFailure +from ._linenos import pos_range from functools import lru_cache @@ -602,15 +603,63 @@ def find_node( position = self.instruction(index).positions assert position is not None and position.lineno is not None - return only( - cast(EnhancedAST, node) - for node in self.source._nodes_by_line[position.lineno] - if isinstance(node, typ) - if not isinstance(node, ast.Expr) - # matchvalue.value has the same positions as matchvalue themself, so we exclude ast.MatchValue - if not isinstance(node, ast.MatchValue) - if all( - getattr(position, attr) == getattr(node, attr) - for attr in match_positions + node = self.source.tree + + # print(ast.dump(node, indent=2, include_attributes=True)) + + def childs(node): + for child in ast.iter_child_nodes(node): + if isinstance(child, typ) and not isinstance(child,ast.Expr): + yield child + else: + yield from childs(child) + + def pos_range(node): + start, end = node, node + if hasattr(node, "decorator_list") and node.decorator_list: + start, end = node.decorator_list[0], node + + return (start.lineno, start.col_offset), ( + end.end_lineno, + end.end_col_offset, ) - ) + + def find(node): + print(node) + if ( + not isinstance(node, (ast.Expr, ast.MatchValue)) + and all( + hasattr(node, attr) + and getattr(position, attr) == getattr(node, attr) + for attr in match_positions + ) + and isinstance(node, typ) + ): + yield node + + for child in childs(node): + print(" child:",child) + if hasattr(child, "lineno"): + start, end = pos_range(child) + if ( + start <= (position.lineno, position.col_offset) + and (position.end_lineno, position.end_col_offset) <= end + ): + yield from find(child) + + node = only(find(self.source.tree)) + + # assert node == only( + # cast(EnhancedAST, node) + # for node in self.source._nodes_by_line[position.lineno] + # if isinstance(node, typ) + # if not isinstance(node, ast.Expr) + # # matchvalue.value has the same positions as matchvalue themself, so we exclude ast.MatchValue + # if not isinstance(node, ast.MatchValue) + # if all( + # getattr(position, attr) == getattr(node, attr) + # for attr in match_positions + # ) + # ) + + return node diff --git a/executing/executing.py b/executing/executing.py index c28091d..d9d857b 100644 --- a/executing/executing.py +++ b/executing/executing.py @@ -39,6 +39,8 @@ from threading import RLock from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, Iterator, List, Optional, Sequence, Set, Sized, Tuple, Type, TypeVar, Union, cast +from ._linenos import LinenosCache + if TYPE_CHECKING: # pragma: no cover from asttokens import ASTTokens, ASTText from asttokens.asttokens import ASTTextBase @@ -193,6 +195,7 @@ def only(it): return lst[0] + class Source(object): """ The source code of a single file and associated metadata. @@ -247,7 +250,9 @@ def __init__(self, filename, lines): for i, line in enumerate(lines) ]) - self._nodes_by_line = defaultdict(list) + if sys.version_info < (3,8): + self._nodes_by_line = defaultdict(list) + self.tree = None self._qualnames = {} self._asttokens = None # type: Optional[ASTTokens] @@ -258,11 +263,17 @@ def __init__(self, filename, lines): except (SyntaxError, ValueError): pass else: + if sys.version_info >= (3,8): + self._nodes_by_line= LinenosCache(self.tree) + for node in ast.walk(self.tree): for child in ast.iter_child_nodes(node): cast(EnhancedAST, child).parent = cast(EnhancedAST, node) - for lineno in node_linenos(node): - self._nodes_by_line[lineno].append(node) + if sys.version_info < (3,8): + for lineno in node_linenos(node): + self._nodes_by_line[lineno].append(node) + + visitor = QualnameVisitor() visitor.visit(self.tree) diff --git a/tests/small_samples/d11c51d03cab4fb5622d7d0331950f63f820940d83005ec86868959f517358fb.py b/tests/small_samples/d11c51d03cab4fb5622d7d0331950f63f820940d83005ec86868959f517358fb.py new file mode 100644 index 0000000..ee1e384 --- /dev/null +++ b/tests/small_samples/d11c51d03cab4fb5622d7d0331950f63f820940d83005ec86868959f517358fb.py @@ -0,0 +1 @@ +f'Customize {self}' \ No newline at end of file