Repository: pylint-dev/astroid

Issue:
Impossible to get correct line number for the docstrings
# Current behavior: 

Removing the first `ast.Expr` from the tree to populate `.doc` property makes it impossible to get correct line number for the docstirng. 

We ca start a module docstring at the line 3:

```
#!/usr/bin/env python3

"""module docstring"""
...
```

Maybe there is a workaround that I did find, but it looks like the constant node information is lost. 

The core of the issue lies down in the `TreeRebuilder._get_doc()` method. 

https://github.com/PyCQA/astroid/blob/82bf77d54393a52ecd07d50a791f5e1a63369f11/astroid/rebuilder.py#L122

# Expected behavior: 

The expected behaviour, in my opinion is not to remove the `Expr` node from the container such that users can use correct line numbers and report docstring warnings with better accuracy. 

Astroid version: 2.9.3


Discussion / hints:

Fix (patch):
diff --git a/ChangeLog b/ChangeLog
index 670628b8da..ea4c1e3319 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -6,6 +6,8 @@ What's New in astroid 2.11.0?
 =============================
 Release date: TBA
 
+* Add new (optional) ``doc_node`` attribute to ``nodes.Module``, ``nodes.ClassDef``,
+  and ``nodes.FunctionDef``.
 
 
 What's New in astroid 2.10.1?
@@ -13,7 +15,6 @@ What's New in astroid 2.10.1?
 Release date: TBA
 
 
-
 What's New in astroid 2.10.0?
 =============================
 Release date: 2022-02-27
diff --git a/astroid/const.py b/astroid/const.py
index 93bf19e13a..74b97cfb7c 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 @@
 
 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 182ec8f4a1..a8dcf3549f 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 @@ def __init__(
         """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 @@ def __init__(
 
     # 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 @@ def __init__(
         :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 @@ def postinit(
         type_comment_args=None,
         *,
         position: Optional[Position] = None,
+        doc_node: Optional[Const] = None,
     ):
         """Do some setup after initialisation.
 
@@ -1589,6 +1599,8 @@ def postinit(
             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 @@ def postinit(
         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 @@ def my_meth(self, arg):
     # 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 @@ def __init__(
         :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 @@ def postinit(
         keywords=None,
         *,
         position: Optional[Position] = None,
+        doc_node: Optional[Const] = None,
     ):
         """Do some setup after initialisation.
 
@@ -2280,6 +2298,8 @@ def postinit(
         :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 @@ def postinit(
         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 7a025e8fac..141494ea72 100644
--- a/astroid/rebuilder.py
+++ b/astroid/rebuilder.py
@@ -33,6 +33,7 @@
 
 import sys
 import token
+import tokenize
 from io import StringIO
 from tokenize import TokenInfo, generate_tokens
 from typing import (
@@ -42,6 +43,7 @@
     Generator,
     List,
     Optional,
+    Set,
     Tuple,
     Type,
     TypeVar,
@@ -52,7 +54,7 @@
 
 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_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 @@ def __init__(
             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 @@ def _get_doc(self, node: T_Doc) -> Tuple[T_Doc, Optional[str]]:
                     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 @@ def _get_position_info(
 
         # 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 @@ def visit_module(
 
         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 @@ def visit_module(
             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 @@ def visit_classdef(
         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 @@ def visit_classdef(
                 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 @@ def _visit_functiondef(
     ) -> 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 @@ def _visit_functiondef(
             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
 


Test changes:
diff --git a/tests/unittest_nodes.py b/tests/unittest_nodes.py
index 1ac896fc02..a7735c5c6e 100644
--- a/tests/unittest_nodes.py
+++ b/tests/unittest_nodes.py
@@ -1596,23 +1596,32 @@ class SomeClass:
 
 
 def test_get_doc() -> None:
-    node = astroid.extract_node(
-        """
+    code = textwrap.dedent(
+        """\
     def func():
         "Docstring"
         return 1
     """
     )
+    node: nodes.FunctionDef = astroid.extract_node(code)  # type: ignore[assignment]
     assert node.doc == "Docstring"
-
-    node = astroid.extract_node(
-        """
+    assert isinstance(node.doc_node, nodes.Const)
+    assert node.doc_node.value == "Docstring"
+    assert node.doc_node.lineno == 2
+    assert node.doc_node.col_offset == 4
+    assert node.doc_node.end_lineno == 2
+    assert node.doc_node.end_col_offset == 15
+
+    code = textwrap.dedent(
+        """\
     def func():
         ...
         return 1
     """
     )
+    node = astroid.extract_node(code)
     assert node.doc is None
+    assert node.doc_node is None
 
 
 @test_utils.require_version(minver="3.8")
diff --git a/tests/unittest_scoped_nodes.py b/tests/unittest_scoped_nodes.py
index 9a156781dc..6a95385cec 100644
--- a/tests/unittest_scoped_nodes.py
+++ b/tests/unittest_scoped_nodes.py
@@ -295,6 +295,69 @@ def test_stream_api(self) -> None:
             with open(path, "rb") as file_io:
                 self.assertEqual(stream.read(), file_io.read())
 
+    @staticmethod
+    def test_singleline_docstring() -> None:
+        data = textwrap.dedent(
+            """\
+            '''Hello World'''
+            foo = 1
+        """
+        )
+        module = builder.parse(data, __name__)
+        assert isinstance(module.doc_node, nodes.Const)
+        assert module.doc_node.lineno == 1
+        assert module.doc_node.col_offset == 0
+        assert module.doc_node.end_lineno == 1
+        assert module.doc_node.end_col_offset == 17
+
+    @staticmethod
+    def test_multiline_docstring() -> None:
+        data = textwrap.dedent(
+            """\
+            '''Hello World
+
+            Also on this line.
+            '''
+            foo = 1
+        """
+        )
+        module = builder.parse(data, __name__)
+
+        assert isinstance(module.doc_node, nodes.Const)
+        assert module.doc_node.lineno == 1
+        assert module.doc_node.col_offset == 0
+        assert module.doc_node.end_lineno == 4
+        assert module.doc_node.end_col_offset == 3
+
+    @staticmethod
+    def test_comment_before_docstring() -> None:
+        data = textwrap.dedent(
+            """\
+            # Some comment
+            '''This is
+
+            a multiline docstring.
+            '''
+        """
+        )
+        module = builder.parse(data, __name__)
+
+        assert isinstance(module.doc_node, nodes.Const)
+        assert module.doc_node.lineno == 2
+        assert module.doc_node.col_offset == 0
+        assert module.doc_node.end_lineno == 5
+        assert module.doc_node.end_col_offset == 3
+
+    @staticmethod
+    def test_without_docstring() -> None:
+        data = textwrap.dedent(
+            """\
+            foo = 1
+        """
+        )
+        module = builder.parse(data, __name__)
+        assert module.doc_node is None
+
 
 class FunctionNodeTest(ModuleLoader, unittest.TestCase):
     def test_special_attributes(self) -> None:
@@ -752,6 +815,118 @@ def test(cls):
         self.assertIsInstance(inferred, nodes.ClassDef)
         self.assertEqual(inferred.name, "MyClass")
 
+    @staticmethod
+    def test_singleline_docstring() -> None:
+        code = textwrap.dedent(
+            """\
+            def foo():
+                '''Hello World'''
+                bar = 1
+        """
+        )
+        func: nodes.FunctionDef = builder.extract_node(code)  # type: ignore[assignment]
+
+        assert isinstance(func.doc_node, nodes.Const)
+        assert func.doc_node.lineno == 2
+        assert func.doc_node.col_offset == 4
+        assert func.doc_node.end_lineno == 2
+        assert func.doc_node.end_col_offset == 21
+
+    @staticmethod
+    def test_multiline_docstring() -> None:
+        code = textwrap.dedent(
+            """\
+            def foo():
+                '''Hello World
+
+                Also on this line.
+                '''
+                bar = 1
+        """
+        )
+        func: nodes.FunctionDef = builder.extract_node(code)  # type: ignore[assignment]
+
+        assert isinstance(func.doc_node, nodes.Const)
+        assert func.doc_node.lineno == 2
+        assert func.doc_node.col_offset == 4
+        assert func.doc_node.end_lineno == 5
+        assert func.doc_node.end_col_offset == 7
+
+    @staticmethod
+    def test_multiline_docstring_async() -> None:
+        code = textwrap.dedent(
+            """\
+            async def foo(var: tuple = ()):
+                '''Hello
+
+                World
+                '''
+        """
+        )
+        func: nodes.FunctionDef = builder.extract_node(code)  # type: ignore[assignment]
+
+        assert isinstance(func.doc_node, nodes.Const)
+        assert func.doc_node.lineno == 2
+        assert func.doc_node.col_offset == 4
+        assert func.doc_node.end_lineno == 5
+        assert func.doc_node.end_col_offset == 7
+
+    @staticmethod
+    def test_docstring_special_cases() -> None:
+        code = textwrap.dedent(
+            """\
+        def f1(var: tuple = ()):  #@
+            'Hello World'
+
+        def f2() -> "just some comment with an open bracket(":  #@
+            'Hello World'
+
+        def f3() -> "Another comment with a colon: ":  #@
+            'Hello World'
+
+        def f4():  #@
+            # It should work with comments too
+            'Hello World'
+        """
+        )
+        ast_nodes: List[nodes.FunctionDef] = builder.extract_node(code)  # type: ignore[assignment]
+        assert len(ast_nodes) == 4
+
+        assert isinstance(ast_nodes[0].doc_node, nodes.Const)
+        assert ast_nodes[0].doc_node.lineno == 2
+        assert ast_nodes[0].doc_node.col_offset == 4
+        assert ast_nodes[0].doc_node.end_lineno == 2
+        assert ast_nodes[0].doc_node.end_col_offset == 17
+
+        assert isinstance(ast_nodes[1].doc_node, nodes.Const)
+        assert ast_nodes[1].doc_node.lineno == 5
+        assert ast_nodes[1].doc_node.col_offset == 4
+        assert ast_nodes[1].doc_node.end_lineno == 5
+        assert ast_nodes[1].doc_node.end_col_offset == 17
+
+        assert isinstance(ast_nodes[2].doc_node, nodes.Const)
+        assert ast_nodes[2].doc_node.lineno == 8
+        assert ast_nodes[2].doc_node.col_offset == 4
+        assert ast_nodes[2].doc_node.end_lineno == 8
+        assert ast_nodes[2].doc_node.end_col_offset == 17
+
+        assert isinstance(ast_nodes[3].doc_node, nodes.Const)
+        assert ast_nodes[3].doc_node.lineno == 12
+        assert ast_nodes[3].doc_node.col_offset == 4
+        assert ast_nodes[3].doc_node.end_lineno == 12
+        assert ast_nodes[3].doc_node.end_col_offset == 17
+
+    @staticmethod
+    def test_without_docstring() -> None:
+        code = textwrap.dedent(
+            """\
+            def foo():
+                bar = 1
+        """
+        )
+        func: nodes.FunctionDef = builder.extract_node(code)  # type: ignore[assignment]
+        assert func.doc_node is None
+
 
 class ClassNodeTest(ModuleLoader, unittest.TestCase):
     def test_dict_interface(self) -> None:
@@ -2088,6 +2263,52 @@ def update(self):
         # Should not crash
         builder.parse(data)
 
+    @staticmethod
+    def test_singleline_docstring() -> None:
+        code = textwrap.dedent(
+            """\
+            class Foo:
+                '''Hello World'''
+                bar = 1
+        """
+        )
+        node: nodes.ClassDef = builder.extract_node(code)  # type: ignore[assignment]
+        assert isinstance(node.doc_node, nodes.Const)
+        assert node.doc_node.lineno == 2
+        assert node.doc_node.col_offset == 4
+        assert node.doc_node.end_lineno == 2
+        assert node.doc_node.end_col_offset == 21
+
+    @staticmethod
+    def test_multiline_docstring() -> None:
+        code = textwrap.dedent(
+            """\
+            class Foo:
+                '''Hello World
+
+                Also on this line.
+                '''
+                bar = 1
+        """
+        )
+        node: nodes.ClassDef = builder.extract_node(code)  # type: ignore[assignment]
+        assert isinstance(node.doc_node, nodes.Const)
+        assert node.doc_node.lineno == 2
+        assert node.doc_node.col_offset == 4
+        assert node.doc_node.end_lineno == 5
+        assert node.doc_node.end_col_offset == 7
+
+    @staticmethod
+    def test_without_docstring() -> None:
+        code = textwrap.dedent(
+            """\
+            class Foo:
+                bar = 1
+        """
+        )
+        node: nodes.ClassDef = builder.extract_node(code)  # type: ignore[assignment]
+        assert node.doc_node is None
+
 
 def test_issue940_metaclass_subclass_property() -> None:
     node = builder.extract_node(