Add position attribute for nodes (#1393)

Author: cdce8p

ChangeLog

@@ -53,6 +53,11 @@ Release date: TBA
 
   Closes #1085
 
+* Add optional ``NodeNG.position`` attribute.
+  Used for block nodes to highlight position of keyword(s) and name
+  in cases where the AST doesn't provide good enough positional information.
+  E.g. ``nodes.ClassDef``, ``nodes.FunctionDef``.
+
 * Fix ``ClassDef.fromlineno``. For Python < 3.8 the ``lineno`` attribute includes decorators.
   ``fromlineno`` should return the line of the ``class`` statement itself.
 

astroid/__init__.py

@@ -43,6 +43,8 @@
 * builder contains the class responsible to build astroid trees
 """
 
+import functools
+import tokenize
 from importlib import import_module
 from pathlib import Path
 
@@ -60,7 +62,7 @@
 from astroid.bases import BaseInstance, BoundMethod, Instance, UnboundMethod
 from astroid.brain.helpers import register_module_extender
 from astroid.builder import extract_node, parse
-from astroid.const import Context, Del, Load, Store
+from astroid.const import PY310_PLUS, Context, Del, Load, Store
 from astroid.exceptions import *
 from astroid.inference_tip import _inference_tip_cached, inference_tip
 from astroid.objects import ExceptionInstance
@@ -165,6 +167,15 @@
 
 from astroid.util import Uninferable
 
+# Performance hack for tokenize. See https://bugs.python.org/issue43014
+# Adapted from https://github.com/PyCQA/pycodestyle/pull/993
+if (
+    not PY310_PLUS
+    and callable(getattr(tokenize, "_compile", None))
+    and getattr(tokenize._compile, "__wrapped__", None) is None  # type: ignore[attr-defined]
+):
+    tokenize._compile = functools.lru_cache()(tokenize._compile)  # type: ignore[attr-defined]
+
 # load brain plugins
 ASTROID_INSTALL_DIRECTORY = Path(__file__).parent
 BRAIN_MODULES_DIRECTORY = ASTROID_INSTALL_DIRECTORY / "brain"

astroid/builder.py

@@ -175,7 +175,7 @@ def _post_build(self, module, encoding):
             module = self._manager.visit_transforms(module)
         return module
 
-    def _data_build(self, data, modname, path):
+    def _data_build(self, data: str, modname, path):
         """Build tree node from data and add some information"""
         try:
             node, parser_module = _parse_string(data, type_comments=True)
@@ -200,7 +200,7 @@ def _data_build(self, data, modname, path):
                 path is not None
                 and os.path.splitext(os.path.basename(path))[0] == "__init__"
             )
-        builder = rebuilder.TreeRebuilder(self._manager, parser_module)
+        builder = rebuilder.TreeRebuilder(self._manager, parser_module, data)
         module = builder.visit_module(node, modname, node_file, package)
         module._import_from_nodes = builder._import_from_nodes
         module._delayed_assattr = builder._delayed_assattr

astroid/nodes/node_ng.py

@@ -28,6 +28,7 @@
 from astroid.manager import AstroidManager
 from astroid.nodes.as_string import AsStringVisitor
 from astroid.nodes.const import OP_PRECEDENCE
+from astroid.nodes.utils import Position
 
 if TYPE_CHECKING:
     from astroid import nodes
@@ -118,6 +119,12 @@ def __init__(
         Note: This is after the last symbol.
         """
 
+        self.position: Optional[Position] = None
+        """Position of keyword(s) and name. Used as fallback for block nodes
+        which might not provide good enough positional information.
+        E.g. ClassDef, FunctionDef.
+        """
+
     def infer(self, context=None, **kwargs):
         """Get a generator of the inferred values.
 

astroid/nodes/scoped_nodes/scoped_nodes.py

@@ -78,6 +78,7 @@
 from astroid.interpreter.objectmodel import ClassModel, FunctionModel, ModuleModel
 from astroid.manager import AstroidManager
 from astroid.nodes import Arguments, Const, node_classes
+from astroid.nodes.utils import Position
 
 if sys.version_info >= (3, 6, 2):
     from typing import NoReturn
@@ -1490,7 +1491,7 @@ class FunctionDef(mixins.MultiLineBlockMixin, node_classes.Statement, Lambda):
     type_comment_returns = None
     """If present, this will contain the return type annotation, passed by a type comment"""
     # attributes below are set by the builder module or by raw factories
-    _other_fields = ("name", "doc")
+    _other_fields = ("name", "doc", "position")
     _other_other_fields = (
         "locals",
         "_type",
@@ -1567,6 +1568,8 @@ def postinit(
         returns=None,
         type_comment_returns=None,
         type_comment_args=None,
+        *,
+        position: Optional[Position] = None,
     ):
         """Do some setup after initialisation.
 
@@ -1582,13 +1585,16 @@ def postinit(
             The return type annotation passed via a type comment.
         :params type_comment_args:
             The args type annotation passed via a type comment.
+        :params position:
+            Position of function keyword(s) and name.
         """
         self.args = args
         self.body = body
         self.decorators = decorators
         self.returns = returns
         self.type_comment_returns = type_comment_returns
         self.type_comment_args = type_comment_args
+        self.position = position
 
     @decorators_mod.cachedproperty
     def extra_decorators(self) -> List[node_classes.Call]:
@@ -2131,7 +2137,7 @@ def my_meth(self, arg):
             ":type: str"
         ),
     )
-    _other_fields = ("name", "doc", "is_dataclass")
+    _other_fields = ("name", "doc", "is_dataclass", "position")
     _other_other_fields = ("locals", "_newstyle")
     _newstyle = None
 
@@ -2241,7 +2247,15 @@ def implicit_locals(self):
 
     # pylint: disable=redefined-outer-name
     def postinit(
-        self, bases, body, decorators, newstyle=None, metaclass=None, keywords=None
+        self,
+        bases,
+        body,
+        decorators,
+        newstyle=None,
+        metaclass=None,
+        keywords=None,
+        *,
+        position: Optional[Position] = None,
     ):
         """Do some setup after initialisation.
 
@@ -2262,6 +2276,8 @@ def postinit(
 
         :param keywords: The keywords given to the class definition.
         :type keywords: list(Keyword) or None
+
+        :param position: Position of class keyword and name.
         """
         if keywords is not None:
             self.keywords = keywords
@@ -2272,6 +2288,7 @@ def postinit(
             self._newstyle = newstyle
         if metaclass is not None:
             self._metaclass = metaclass
+        self.position = position
 
     def _newstyle_impl(self, context=None):
         if context is None:

astroid/nodes/utils.py

@@ -0,0 +1,10 @@
+from typing import NamedTuple
+
+
+class Position(NamedTuple):
+    """Position with line and column information."""
+
+    lineno: int
+    col_offset: int
+    end_lineno: int
+    end_col_offset: int

astroid/rebuilder.py

@@ -31,6 +31,9 @@
 """
 
 import sys
+import token
+from io import StringIO
+from tokenize import TokenInfo, generate_tokens
 from typing import (
     TYPE_CHECKING,
     Callable,
@@ -48,9 +51,10 @@
 
 from astroid import nodes
 from astroid._ast import ParserModule, get_parser_module, parse_function_type_comment
-from astroid.const import PY38, PY38_PLUS, Context
+from astroid.const import PY36, PY38, PY38_PLUS, Context
 from astroid.manager import AstroidManager
 from astroid.nodes import NodeNG
+from astroid.nodes.utils import Position
 
 if sys.version_info >= (3, 8):
     from typing import Final
@@ -88,9 +92,13 @@ class TreeRebuilder:
     """Rebuilds the _ast tree to become an Astroid tree"""
 
     def __init__(
-        self, manager: AstroidManager, parser_module: Optional[ParserModule] = None
-    ):
+        self,
+        manager: AstroidManager,
+        parser_module: Optional[ParserModule] = None,
+        data: Optional[str] = None,
+    ) -> None:
         self._manager = manager
+        self._data = data.split("\n") if data else None
         self._global_names: List[Dict[str, List[nodes.Global]]] = []
         self._import_from_nodes: List[nodes.ImportFrom] = []
         self._delayed_assattr: List[nodes.AssignAttr] = []
@@ -133,6 +141,68 @@ def _get_context(
     ) -> Context:
         return self._parser_module.context_classes.get(type(node.ctx), Context.Load)
 
+    def _get_position_info(
+        self,
+        node: Union["ast.ClassDef", "ast.FunctionDef", "ast.AsyncFunctionDef"],
+        parent: Union[nodes.ClassDef, nodes.FunctionDef, nodes.AsyncFunctionDef],
+    ) -> Optional[Position]:
+        """Return position information for ClassDef and FunctionDef nodes.
+
+        In contrast to AST positions, these only include the actual keyword(s)
+        and the class / function name.
+
+        >>> @decorator
+        >>> async def some_func(var: int) -> None:
+        >>> ^^^^^^^^^^^^^^^^^^^
+        """
+        if not self._data:
+            return None
+        end_lineno: Optional[int] = getattr(node, "end_lineno", None)
+        if node.body:
+            end_lineno = node.body[0].lineno
+        # pylint: disable-next=unsubscriptable-object
+        data = "\n".join(self._data[node.lineno - 1 : end_lineno])
+
+        start_token: Optional[TokenInfo] = None
+        keyword_tokens: Tuple[int, ...] = (token.NAME,)
+        if isinstance(parent, nodes.AsyncFunctionDef):
+            search_token = "async"
+            if PY36:
+                # In Python 3.6, the token type for 'async' was 'ASYNC'
+                # In Python 3.7, the type was changed to 'NAME' and 'ASYNC' removed
+                # Python 3.8 added it back. However, if we use it unconditionally
+                # we would break 3.7.
+                keyword_tokens = (token.NAME, token.ASYNC)
+        elif isinstance(parent, nodes.FunctionDef):
+            search_token = "def"
+        else:
+            search_token = "class"
+
+        for t in generate_tokens(StringIO(data).readline):
+            if (
+                start_token is not None
+                and t.type == token.NAME
+                and t.string == node.name
+            ):
+                break
+            if t.type in keyword_tokens:
+                if t.string == search_token:
+                    start_token = t
+                    continue
+                if t.string in {"def"}:
+                    continue
+            start_token = None
+        else:
+            return None
+
+        # pylint: disable=undefined-loop-variable
+        return Position(
+            lineno=node.lineno - 1 + start_token.start[0],
+            col_offset=start_token.start[1],
+            end_lineno=node.lineno - 1 + t.end[0],
+            end_col_offset=t.end[1],
+        )
+
     def visit_module(
         self, node: "ast.Module", modname: str, modpath: str, package: bool
     ) -> nodes.Module:
@@ -1203,6 +1273,7 @@ def visit_classdef(
                 for kwd in node.keywords
                 if kwd.arg != "metaclass"
             ],
+            position=self._get_position_info(node, newnode),
         )
         return newnode
 
@@ -1551,6 +1622,7 @@ def _visit_functiondef(
             returns=returns,
             type_comment_returns=type_comment_returns,
             type_comment_args=type_comment_args,
+            position=self._get_position_info(node, newnode),
         )
         self._global_names.pop()
         return newnode