Pass doc_node to postinit of all PartialFunction constructors (#1437)

Co-authored-by: Marc Mueller <[email protected]>

Author: DanielNoord

astroid/brain/brain_functools.py

@@ -7,8 +7,10 @@
 """Astroid hooks for understanding functools library module."""
 from functools import partial
 from itertools import chain
+from typing import Iterator, Optional
 
-from astroid import BoundMethod, arguments, extract_node, helpers, objects
+from astroid import BoundMethod, arguments, extract_node, helpers, nodes, objects
+from astroid.context import InferenceContext
 from astroid.exceptions import InferenceError, UseInferenceDefault
 from astroid.inference_tip import inference_tip
 from astroid.interpreter import objectmodel
@@ -62,7 +64,9 @@ def _transform_lru_cache(node, context=None) -> None:
     node.special_attributes = LruWrappedModel()(node)
 
 
-def _functools_partial_inference(node, context=None):
+def _functools_partial_inference(
+    node: nodes.Call, context: Optional[InferenceContext] = None
+) -> Iterator[objects.PartialFunction]:
     call = arguments.CallSite.from_call(node, context=context)
     number_of_positional = len(call.positional_arguments)
     if number_of_positional < 1:
@@ -101,7 +105,6 @@ def _functools_partial_inference(node, context=None):
     partial_function = objects.PartialFunction(
         call,
         name=inferred_wrapped_function.name,
-        doc=inferred_wrapped_function.doc,
         lineno=inferred_wrapped_function.lineno,
         col_offset=inferred_wrapped_function.col_offset,
         parent=node.parent,
@@ -113,6 +116,7 @@ def _functools_partial_inference(node, context=None):
         returns=inferred_wrapped_function.returns,
         type_comment_returns=inferred_wrapped_function.type_comment_returns,
         type_comment_args=inferred_wrapped_function.type_comment_args,
+        doc_node=inferred_wrapped_function.doc_node,
     )
     return iter((partial_function,))
 

astroid/objects.py

@@ -269,6 +269,7 @@ class PartialFunction(scoped_nodes.FunctionDef):
     def __init__(
         self, call, name=None, doc=None, lineno=None, col_offset=None, parent=None
     ):
+        # TODO: Pass end_lineno and end_col_offset as well
         super().__init__(name, doc, lineno, col_offset, parent=None)
         # A typical FunctionDef automatically adds its name to the parent scope,
         # but a partial should not, so defer setting parent until after init

tests/unittest_brain.py

@@ -2994,6 +2994,30 @@ def test_infer_dict_from_keys() -> None:
 
 
 class TestFunctoolsPartial:
+    @staticmethod
+    def test_infer_partial() -> None:
+        ast_node = astroid.extract_node(
+            """
+        from functools import partial
+        def test(a, b):
+            '''Docstring'''
+            return a + b
+        partial(test, 1)(3) #@
+        """
+        )
+        assert isinstance(ast_node.func, nodes.Call)
+        inferred = ast_node.func.inferred()
+        assert len(inferred) == 1
+        partial = inferred[0]
+        assert isinstance(partial, objects.PartialFunction)
+        assert isinstance(partial.doc_node, nodes.Const)
+        assert partial.doc_node.value == "Docstring"
+        with pytest.warns(DeprecationWarning) as records:
+            assert partial.doc == "Docstring"
+            assert len(records) == 1
+        assert partial.lineno == 3
+        assert partial.col_offset == 0
+
     def test_invalid_functools_partial_calls(self) -> None:
         ast_nodes = astroid.extract_node(
             """