Add doc_node attribute (#1276)

Co-authored-by: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Co-authored-by: Daniël van Noord <13665637+DanielNoord@users.noreply.github.com>

3 files changed, 113 insertions, 14 deletions

diff --git a/astroid/const.py b/astroid/const.py
index 93bf19e1..74b97cfb 100644
--- a/astroid/const.py
+++ b/astroid/const.py
@@ -1,4 +1,5 @@
 import enum
+import platform
 import sys
 
 PY36 = sys.version_info[:2] == (3, 6)
@@ -11,6 +12,8 @@ BUILTINS = "builtins"  # TODO Remove in 2.8
 
 WIN32 = sys.platform == "win32"
 
+IMPLEMENTATION_PYPY = platform.python_implementation() == "PyPy"
+
 
 class Context(enum.Enum):
     Load = 1
diff --git a/astroid/nodes/scoped_nodes/scoped_nodes.py b/astroid/nodes/scoped_nodes/scoped_nodes.py
index 182ec8f4..a8dcf354 100644
--- a/astroid/nodes/scoped_nodes/scoped_nodes.py
+++ b/astroid/nodes/scoped_nodes/scoped_nodes.py
@@ -386,7 +386,7 @@ class Module(LocalsDictNodeNG):
     <Module l.0 at 0x7f23b2e4eda0>
     """
 
-    _astroid_fields = ("body",)
+    _astroid_fields = ("doc_node", "body")
 
     fromlineno: Literal[0] = 0
     """The first line that this node appears on in the source code."""
@@ -479,10 +479,14 @@ class Module(LocalsDictNodeNG):
         """A map of the name of a global variable to the node defining the global."""
 
         self.locals = self.globals = {}
+        """A map of the name of a local variable to the node defining the local."""
 
         self.body: Optional[List[node_classes.NodeNG]] = []
         """The contents of the module."""
 
+        self.doc_node: Optional[Const] = None
+        """The doc node associated with this node."""
+
         self.future_imports: Set[str] = set()
         """The imports from ``__future__``."""
 
@@ -490,13 +494,15 @@ class Module(LocalsDictNodeNG):
 
     # pylint: enable=redefined-builtin
 
-    def postinit(self, body=None):
+    def postinit(self, body=None, *, doc_node: Optional[Const] = None):
         """Do some setup after initialisation.
 
         :param body: The contents of the module.
         :type body: list(NodeNG) or None
+        :param doc_node: The doc node associated with this node.
         """
         self.body = body
+        self.doc_node = doc_node
 
     def _get_stream(self):
         if self.file_bytes is not None:
@@ -1463,7 +1469,7 @@ class FunctionDef(mixins.MultiLineBlockMixin, node_classes.Statement, Lambda):
     <FunctionDef.my_func l.2 at 0x7f23b2e71e10>
     """
 
-    _astroid_fields = ("decorators", "args", "returns", "body")
+    _astroid_fields = ("decorators", "args", "returns", "doc_node", "body")
     _multi_line_block_fields = ("body",)
     returns = None
     decorators: Optional[node_classes.Decorators] = None
@@ -1549,6 +1555,9 @@ class FunctionDef(mixins.MultiLineBlockMixin, node_classes.Statement, Lambda):
         :type doc: str or None
         """
 
+        self.doc_node: Optional[Const] = None
+        """The doc node associated with this node."""
+
         self.instance_attrs = {}
         super().__init__(
             lineno=lineno,
@@ -1572,6 +1581,7 @@ class FunctionDef(mixins.MultiLineBlockMixin, node_classes.Statement, Lambda):
         type_comment_args=None,
         *,
         position: Optional[Position] = None,
+        doc_node: Optional[Const] = None,
     ):
         """Do some setup after initialisation.
 
@@ -1589,6 +1599,8 @@ class FunctionDef(mixins.MultiLineBlockMixin, node_classes.Statement, Lambda):
             The args type annotation passed via a type comment.
         :params position:
             Position of function keyword(s) and name.
+        :param doc_node:
+            The doc node associated with this node.
         """
         self.args = args
         self.body = body
@@ -1597,6 +1609,7 @@ class FunctionDef(mixins.MultiLineBlockMixin, node_classes.Statement, Lambda):
         self.type_comment_returns = type_comment_returns
         self.type_comment_args = type_comment_args
         self.position = position
+        self.doc_node = doc_node
 
     @decorators_mod.cachedproperty
     def extra_decorators(self) -> List[node_classes.Call]:
@@ -2098,6 +2111,7 @@ def get_wrapping_class(node):
     return klass
 
 
+# pylint: disable=too-many-instance-attributes
 class ClassDef(mixins.FilterStmtsMixin, LocalsDictNodeNG, node_classes.Statement):
     """Class representing an :class:`ast.ClassDef` node.
 
@@ -2115,7 +2129,7 @@ class ClassDef(mixins.FilterStmtsMixin, LocalsDictNodeNG, node_classes.Statement
     # by a raw factories
 
     # a dictionary of class instances attributes
-    _astroid_fields = ("decorators", "bases", "keywords", "body")  # name
+    _astroid_fields = ("decorators", "bases", "keywords", "doc_node", "body")  # name
 
     decorators = None
     """The decorators that are applied to this class.
@@ -2217,6 +2231,9 @@ class ClassDef(mixins.FilterStmtsMixin, LocalsDictNodeNG, node_classes.Statement
         :type doc: str or None
         """
 
+        self.doc_node: Optional[Const] = None
+        """The doc node associated with this node."""
+
         self.is_dataclass: bool = False
         """Whether this class is a dataclass."""
 
@@ -2258,6 +2275,7 @@ class ClassDef(mixins.FilterStmtsMixin, LocalsDictNodeNG, node_classes.Statement
         keywords=None,
         *,
         position: Optional[Position] = None,
+        doc_node: Optional[Const] = None,
     ):
         """Do some setup after initialisation.
 
@@ -2280,6 +2298,8 @@ class ClassDef(mixins.FilterStmtsMixin, LocalsDictNodeNG, node_classes.Statement
         :type keywords: list(Keyword) or None
 
         :param position: Position of class keyword and name.
+
+        :param doc_node: The doc node associated with this node.
         """
         if keywords is not None:
             self.keywords = keywords
@@ -2291,6 +2311,7 @@ class ClassDef(mixins.FilterStmtsMixin, LocalsDictNodeNG, node_classes.Statement
         if metaclass is not None:
             self._metaclass = metaclass
         self.position = position
+        self.doc_node = doc_node
 
     def _newstyle_impl(self, context=None):
         if context is None:
diff --git a/astroid/rebuilder.py b/astroid/rebuilder.py
index 7a025e8f..141494ea 100644
--- a/astroid/rebuilder.py
+++ b/astroid/rebuilder.py
@@ -33,6 +33,7 @@ order to get a single Astroid representation
 
 import sys
 import token
+import tokenize
 from io import StringIO
 from tokenize import TokenInfo, generate_tokens
 from typing import (
@@ -42,6 +43,7 @@ from typing import (
     Generator,
     List,
     Optional,
+    Set,
     Tuple,
     Type,
     TypeVar,
@@ -52,7 +54,7 @@ from typing import (
 
 from astroid import nodes
 from astroid._ast import ParserModule, get_parser_module, parse_function_type_comment
-from astroid.const import PY36, PY38, PY38_PLUS, Context
+from astroid.const import IMPLEMENTATION_PYPY, PY36, PY38, PY38_PLUS, Context
 from astroid.manager import AstroidManager
 from astroid.nodes import NodeNG
 from astroid.nodes.utils import Position
@@ -86,6 +88,7 @@ T_Doc = TypeVar(
 T_Function = TypeVar("T_Function", nodes.FunctionDef, nodes.AsyncFunctionDef)
 T_For = TypeVar("T_For", nodes.For, nodes.AsyncFor)
 T_With = TypeVar("T_With", nodes.With, nodes.AsyncWith)
+NodesWithDocsType = Union[nodes.Module, nodes.ClassDef, nodes.FunctionDef]
 
 
 # noinspection PyMethodMayBeStatic
@@ -113,7 +116,10 @@ class TreeRebuilder:
             self._parser_module = parser_module
         self._module = self._parser_module.module
 
-    def _get_doc(self, node: T_Doc) -> Tuple[T_Doc, Optional[str]]:
+    def _get_doc(
+        self, node: T_Doc
+    ) -> Tuple[T_Doc, Optional["ast.Constant | ast.Str"], Optional[str]]:
+        """Return the doc ast node and the actual docstring."""
         try:
             if node.body and isinstance(node.body[0], self._module.Expr):
                 first_value = node.body[0].value
@@ -122,12 +128,17 @@ class TreeRebuilder:
                     and isinstance(first_value, self._module.Constant)
                     and isinstance(first_value.value, str)
                 ):
+                    doc_ast_node = first_value
                     doc = first_value.value if PY38_PLUS else first_value.s
                     node.body = node.body[1:]
-                    return node, doc
+                    # The ast parser of python < 3.8 sets col_offset of multi-line strings to -1
+                    # as it is unable to determine the value correctly. We reset this to None.
+                    if doc_ast_node.col_offset == -1:
+                        doc_ast_node.col_offset = None
+                    return node, doc_ast_node, doc
         except IndexError:
             pass  # ast built from scratch
-        return node, None
+        return node, None, None
 
     def _get_context(
         self,
@@ -198,12 +209,68 @@ class TreeRebuilder:
 
         # pylint: disable=undefined-loop-variable
         return Position(
-            lineno=node.lineno - 1 + start_token.start[0],
+            lineno=node.lineno + start_token.start[0] - 1,
             col_offset=start_token.start[1],
-            end_lineno=node.lineno - 1 + t.end[0],
+            end_lineno=node.lineno + t.end[0] - 1,
             end_col_offset=t.end[1],
         )
 
+    def _fix_doc_node_position(self, node: NodesWithDocsType) -> None:
+        """Fix start and end position of doc nodes for Python < 3.8."""
+        if not self._data or not node.doc_node or node.lineno is None:
+            return
+        if PY38_PLUS:
+            return
+
+        lineno = node.lineno or 1  # lineno of modules is 0
+        end_range: Optional[int] = node.doc_node.lineno
+        if IMPLEMENTATION_PYPY:
+            end_range = None
+        # pylint: disable-next=unsubscriptable-object
+        data = "\n".join(self._data[lineno - 1 : end_range])
+
+        found_start, found_end = False, False
+        open_brackets = 0
+        skip_token: Set[int] = {token.NEWLINE, token.INDENT}
+        if PY36:
+            skip_token.update((tokenize.NL, tokenize.COMMENT))
+        else:
+            # token.NL and token.COMMENT were added in 3.7
+            skip_token.update((token.NL, token.COMMENT))
+
+        if isinstance(node, nodes.Module):
+            found_end = True
+
+        for t in generate_tokens(StringIO(data).readline):
+            if found_end is False:
+                if (
+                    found_start is False
+                    and t.type == token.NAME
+                    and t.string in {"def", "class"}
+                ):
+                    found_start = True
+                elif found_start is True and t.type == token.OP:
+                    if t.exact_type == token.COLON and open_brackets == 0:
+                        found_end = True
+                    elif t.exact_type == token.LPAR:
+                        open_brackets += 1
+                    elif t.exact_type == token.RPAR:
+                        open_brackets -= 1
+                continue
+            if t.type in skip_token:
+                continue
+            if t.type == token.STRING:
+                break
+            return
+        else:
+            return
+
+        # pylint: disable=undefined-loop-variable
+        node.doc_node.lineno = lineno + t.start[0] - 1
+        node.doc_node.col_offset = t.start[1]
+        node.doc_node.end_lineno = lineno + t.end[0] - 1
+        node.doc_node.end_col_offset = t.end[1]
+
     def visit_module(
         self, node: "ast.Module", modname: str, modpath: str, package: bool
     ) -> nodes.Module:
@@ -211,7 +278,7 @@ class TreeRebuilder:
 
         Note: Method not called by 'visit'
         """
-        node, doc = self._get_doc(node)
+        node, doc_ast_node, doc = self._get_doc(node)
         newnode = nodes.Module(
             name=modname,
             doc=doc,
@@ -220,7 +287,11 @@ class TreeRebuilder:
             package=package,
             parent=None,
         )
-        newnode.postinit([self.visit(child, newnode) for child in node.body])
+        newnode.postinit(
+            [self.visit(child, newnode) for child in node.body],
+            doc_node=self.visit(doc_ast_node, newnode),
+        )
+        self._fix_doc_node_position(newnode)
         return newnode
 
     if sys.version_info >= (3, 10):
@@ -1242,7 +1313,7 @@ class TreeRebuilder:
         self, node: "ast.ClassDef", parent: NodeNG, newstyle: bool = True
     ) -> nodes.ClassDef:
         """visit a ClassDef node to become astroid"""
-        node, doc = self._get_doc(node)
+        node, doc_ast_node, doc = self._get_doc(node)
         if sys.version_info >= (3, 8):
             newnode = nodes.ClassDef(
                 name=node.name,
@@ -1275,7 +1346,9 @@ class TreeRebuilder:
                 if kwd.arg != "metaclass"
             ],
             position=self._get_position_info(node, newnode),
+            doc_node=self.visit(doc_ast_node, newnode),
         )
+        self._fix_doc_node_position(newnode)
         return newnode
 
     def visit_continue(self, node: "ast.Continue", parent: NodeNG) -> nodes.Continue:
@@ -1580,7 +1653,7 @@ class TreeRebuilder:
     ) -> T_Function:
         """visit an FunctionDef node to become astroid"""
         self._global_names.append({})
-        node, doc = self._get_doc(node)
+        node, doc_ast_node, doc = self._get_doc(node)
 
         lineno = node.lineno
         if PY38_PLUS and node.decorator_list:
@@ -1624,7 +1697,9 @@ class TreeRebuilder:
             type_comment_returns=type_comment_returns,
             type_comment_args=type_comment_args,
             position=self._get_position_info(node, newnode),
+            doc_node=self.visit(doc_ast_node, newnode),
         )
+        self._fix_doc_node_position(newnode)
         self._global_names.pop()
         return newnode