Set function __doc__ even when created via constructor

Author: msimacek

graalpython/com.oracle.graal.python.test/src/tests/test_functions.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2018, 2023, Oracle and/or its affiliates. All rights reserved.
+# Copyright (c) 2018, 2024, Oracle and/or its affiliates. All rights reserved.
 # DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 #
 # The Universal Permissive License (UPL), Version 1.0
@@ -278,4 +278,20 @@ def cached_func():
 
     type_entity = getattr(sys.modules['functools'], '_lru_cache_wrapper')
     assert isinstance(cached_func, type_entity), "lru_cache should not be using the python-based version"
-    
+
+
+def test_no_docstring():
+    def no_doc():
+        pass
+
+    assert no_doc.__doc__ is None
+
+
+def test_docstring_via_type_contructor():
+    def foo():
+        """my doc"""
+        return 1
+
+    assert foo.__doc__ == 'my doc'
+    foo2 = type(foo)(foo.__code__, foo.__globals__, foo.__name__, foo.__defaults__, foo.__closure__)
+    assert foo2.__doc__ == 'my doc'

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/function/FunctionBuiltins.java

@@ -28,6 +28,7 @@
 
 import static com.oracle.graal.python.nodes.SpecialAttributeNames.J___CODE__;
 import static com.oracle.graal.python.nodes.SpecialAttributeNames.J___DEFAULTS__;
+import static com.oracle.graal.python.nodes.SpecialAttributeNames.J___DOC__;
 import static com.oracle.graal.python.nodes.SpecialAttributeNames.J___KWDEFAULTS__;
 import static com.oracle.graal.python.nodes.SpecialAttributeNames.J___NAME__;
 import static com.oracle.graal.python.nodes.SpecialAttributeNames.J___QUALNAME__;
@@ -303,4 +304,19 @@ static Object setCode(PFunction self, PCode code,
             return PNone.NONE;
         }
     }
+
+    @Builtin(name = J___DOC__, minNumOfPositionalArgs = 1, maxNumOfPositionalArgs = 2, isGetter = true, isSetter = true)
+    @GenerateNodeFactory
+    abstract static class DocNode extends PythonBinaryBuiltinNode {
+        @Specialization(guards = "isNoValue(none)")
+        Object get(PFunction self, @SuppressWarnings("unused") Object none) {
+            return self.getDoc();
+        }
+
+        @Specialization(guards = "!isNoValue(value)")
+        Object set(PFunction self, Object value) {
+            self.setDoc(value);
+            return PNone.NONE;
+        }
+    }
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/function/PFunction.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017, 2023, Oracle and/or its affiliates.
+ * Copyright (c) 2017, 2024, Oracle and/or its affiliates.
  * Copyright (c) 2013, Regents of the University of California
  *
  * All rights reserved.
@@ -27,10 +27,13 @@
 
 import com.oracle.graal.python.PythonLanguage;
 import com.oracle.graal.python.builtins.PythonBuiltinClassType;
+import com.oracle.graal.python.builtins.objects.PNone;
 import com.oracle.graal.python.builtins.objects.cell.PCell;
 import com.oracle.graal.python.builtins.objects.code.CodeNodes.GetCodeCallTargetNode;
 import com.oracle.graal.python.builtins.objects.code.PCode;
 import com.oracle.graal.python.builtins.objects.object.PythonObject;
+import com.oracle.graal.python.compiler.CodeUnit;
+import com.oracle.graal.python.lib.PyUnicodeCheckNode;
 import com.oracle.graal.python.nodes.PRootNode;
 import com.oracle.graal.python.nodes.builtins.FunctionNodes.GetCallTargetNode;
 import com.oracle.graal.python.runtime.GilNode;
@@ -69,6 +72,7 @@ public final class PFunction extends PythonObject {
     private Object[] defaultValues;
     @CompilationFinal(dimensions = 1) private PKeyword[] finalKwDefaultValues;
     private PKeyword[] kwDefaultValues;
+    private Object doc;
 
     public PFunction(PythonLanguage lang, TruffleString name, TruffleString qualname, PCode code, PythonObject globals, PCell[] closure) {
         this(lang, name, qualname, code, globals, PythonUtils.EMPTY_OBJECT_ARRAY, PKeyword.EMPTY_KEYWORDS, closure);
@@ -145,6 +149,27 @@ public boolean forceSplitDirectCalls() {
         return forceSplitDirectCalls;
     }
 
+    public Object getDoc() {
+        if (CompilerDirectives.injectBranchProbability(CompilerDirectives.SLOWPATH_PROBABILITY, doc == null)) {
+            extractDoc();
+        }
+        return doc;
+    }
+
+    public void setDoc(Object doc) {
+        this.doc = doc;
+    }
+
+    @TruffleBoundary
+    private void extractDoc() {
+        CodeUnit co = getCode().getCodeUnit();
+        if (co != null && co.constants.length > 0 && PyUnicodeCheckNode.executeUncached(co.constants[0])) {
+            doc = co.constants[0];
+        } else {
+            doc = PNone.NONE;
+        }
+    }
+
     @Override
     public final String toString() {
         CompilerAsserts.neverPartOfCompilation();

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/bytecode/MakeFunctionNode.java

@@ -41,7 +41,6 @@
 package com.oracle.graal.python.nodes.bytecode;
 
 import static com.oracle.graal.python.nodes.SpecialAttributeNames.T___ANNOTATIONS__;
-import static com.oracle.graal.python.nodes.SpecialAttributeNames.T___DOC__;
 
 import com.oracle.graal.python.PythonLanguage;
 import com.oracle.graal.python.builtins.objects.cell.PCell;
@@ -64,25 +63,22 @@
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.source.Source;
-import com.oracle.truffle.api.strings.TruffleString;
 
 public abstract class MakeFunctionNode extends PNodeWithContext {
     private final RootCallTarget callTarget;
     private final CodeUnit code;
     private final Signature signature;
-    private final TruffleString doc;
     @CompilationFinal private PCode cachedCode;
 
     private final Assumption sharedCodeStableAssumption = Truffle.getRuntime().createAssumption("shared code stable assumption");
     private final Assumption sharedDefaultsStableAssumption = Truffle.getRuntime().createAssumption("shared defaults stable assumption");
 
     public abstract int execute(VirtualFrame frame, Object globals, int initialStackTop, int flags);
 
-    public MakeFunctionNode(RootCallTarget callTarget, CodeUnit code, Signature signature, TruffleString doc) {
+    public MakeFunctionNode(RootCallTarget callTarget, CodeUnit code, Signature signature) {
         this.callTarget = callTarget;
         this.code = code;
         this.signature = signature;
-        this.doc = doc;
     }
 
     @Specialization
@@ -142,9 +138,6 @@ int makeFunction(VirtualFrame frame, Object globals, int initialStackTop, int fl
         if (annotations != null) {
             writeAttrNode.execute(function, T___ANNOTATIONS__, annotations);
         }
-        if (doc != null) {
-            writeAttrNode.execute(function, T___DOC__, doc);
-        }
 
         frame.setObject(++stackTop, function);
         return stackTop;
@@ -159,11 +152,7 @@ public static MakeFunctionNode create(PythonLanguage language, CodeUnit code, So
         } else {
             callTarget = bytecodeRootNode.getCallTarget();
         }
-        TruffleString doc = null;
-        if (code.constants.length > 0 && code.constants[0] instanceof TruffleString) {
-            doc = (TruffleString) code.constants[0];
-        }
-        return MakeFunctionNodeGen.create(callTarget, code, bytecodeRootNode.getSignature(), doc);
+        return MakeFunctionNodeGen.create(callTarget, code, bytecodeRootNode.getSignature());
     }
 
     public RootCallTarget getCallTarget() {