Create AssignName nodes for ClassDef and FunctionDef

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/scoped_nodes/scoped_nodes.py

@@ -1546,6 +1546,8 @@ def __init__(
         :type doc: str or None
         """
 
+        self.name_node: Optional[node_classes.AssignName] = None
+
         self.instance_attrs = {}
         super().__init__(
             lineno=lineno,
@@ -1567,6 +1569,8 @@ def postinit(
         returns=None,
         type_comment_returns=None,
         type_comment_args=None,
+        *,
+        name_node: Optional[node_classes.AssignName] = None,
     ):
         """Do some setup after initialisation.
 
@@ -1589,6 +1593,7 @@ def postinit(
         self.returns = returns
         self.type_comment_returns = type_comment_returns
         self.type_comment_args = type_comment_args
+        self.name_node = name_node
 
     @decorators_mod.cachedproperty
     def extra_decorators(self) -> List[node_classes.Call]:
@@ -2212,6 +2217,8 @@ def __init__(
         :type doc: str or None
         """
 
+        self.name_node: Optional[node_classes.AssignName] = None
+
         super().__init__(
             lineno=lineno,
             col_offset=col_offset,
@@ -2241,7 +2248,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,
+        *,
+        name_node: Optional[node_classes.AssignName] = None,
     ):
         """Do some setup after initialisation.
 
@@ -2272,6 +2287,7 @@ def postinit(
             self._newstyle = newstyle
         if metaclass is not None:
             self._metaclass = metaclass
+        self.name_node = name_node
 
     def _newstyle_impl(self, context=None):
         if context is None:

astroid/rebuilder.py

@@ -31,6 +31,9 @@
 """
 
 import sys
+import token
+from io import StringIO
+from tokenize import generate_tokens
 from typing import (
     TYPE_CHECKING,
     Callable,
@@ -88,9 +91,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,
     ):
         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 +140,45 @@ def _get_context(
     ) -> Context:
         return self._parser_module.context_classes.get(type(node.ctx), Context.Load)
 
+    def _create_name_node(
+        self,
+        node: Union["ast.ClassDef", "ast.FunctionDef", "ast.AsyncFunctionDef"],
+        parent: Union[nodes.ClassDef, nodes.FunctionDef, nodes.AsyncFunctionDef],
+    ) -> Optional[nodes.AssignName]:
+        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])
+
+        if isinstance(parent, nodes.FunctionDef):
+            search_token_name = "def"
+        else:
+            search_token_name = "class"
+        token_found: bool = False
+
+        for t in generate_tokens(StringIO(data).readline):
+            if token_found and t.type == token.NAME and t.string == node.name:
+                break
+            if t.type == token.NAME and t.string == search_token_name:
+                token_found = True
+                continue
+            token_found = False
+        else:
+            return None
+
+        # pylint: disable=undefined-loop-variable
+        return nodes.AssignName(
+            name=t.string,
+            lineno=node.lineno - 1 + t.start[0],
+            col_offset=t.start[1],
+            end_lineno=node.lineno - 1 + t.end[0],
+            end_col_offset=t.end[1],
+            parent=parent,
+        )
+
     def visit_module(
         self, node: "ast.Module", modname: str, modpath: str, package: bool
     ) -> nodes.Module:
@@ -1203,6 +1249,7 @@ def visit_classdef(
                 for kwd in node.keywords
                 if kwd.arg != "metaclass"
             ],
+            name_node=self._create_name_node(node, newnode),
         )
         return newnode
 
@@ -1551,6 +1598,7 @@ def _visit_functiondef(
             returns=returns,
             type_comment_returns=type_comment_returns,
             type_comment_args=type_comment_args,
+            name_node=self._create_name_node(node, newnode),
         )
         self._global_names.pop()
         return newnode