add support for reading and setting the module docstring

- relevant unittest

Author: cosminbasca

graalpython/com.oracle.graal.python.test/src/tests/package/__init__.py

@@ -37,5 +37,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+"""PACKAGE DOC"""
+
 from . import moduleY
 print("after importing moduleY")

graalpython/com.oracle.graal.python.test/src/tests/package/moduleA.py

@@ -37,4 +37,6 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+"""MODULE A DOC"""
+
 print("module A")

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

@@ -51,6 +51,13 @@ def test_relative_import():
         raise e
 
 
+def test_module_docstring():
+    import package
+    assert package.__doc__ == "PACKAGE DOC"
+    from package import moduleA
+    assert moduleA.__doc__ == "MODULE A DOC"
+
+
 def test_dotted_import():
     # this is to prevent ides from optimising out the unused import
     try:

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/ModuleRootNode.java

@@ -27,27 +27,45 @@
 
 import com.oracle.graal.python.PythonLanguage;
 import com.oracle.graal.python.nodes.expression.ExpressionNode;
+import com.oracle.graal.python.nodes.frame.WriteGlobalNode;
 import com.oracle.truffle.api.CompilerAsserts;
+import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.frame.FrameDescriptor;
 import com.oracle.truffle.api.frame.FrameSlot;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.source.SourceSection;
 
+import static com.oracle.graal.python.nodes.SpecialAttributeNames.__DOC__;
+
 public class ModuleRootNode extends PClosureRootNode {
 
     private final String name;
+    private final String doc;
 
     @Child private ExpressionNode body;
+    @Child private WriteGlobalNode writeModuleDoc;
 
-    public ModuleRootNode(PythonLanguage language, String name, ExpressionNode file, FrameDescriptor descriptor, FrameSlot[] freeVarSlots) {
+    public ModuleRootNode(PythonLanguage language, String name, String doc, ExpressionNode file, FrameDescriptor descriptor, FrameSlot[] freeVarSlots) {
         super(language, descriptor, freeVarSlots);
         this.name = "<module '" + name + "'>";
+        this.doc = doc;
         this.body = file;
     }
 
+    private WriteGlobalNode getWriteModuleDoc() {
+        if (writeModuleDoc == null) {
+            CompilerDirectives.transferToInterpreterAndInvalidate();
+            writeModuleDoc = insert(WriteGlobalNode.create(__DOC__));
+        }
+        return writeModuleDoc;
+    }
+
     @Override
     public Object execute(VirtualFrame frame) {
         addClosureCellsToLocals(frame);
+        if (doc != null) {
+            getWriteModuleDoc().doWrite(frame, doc);
+        }
         return body.execute(frame);
     }
 
@@ -66,6 +84,10 @@ public String getName() {
         return name;
     }
 
+    public String getDoc() {
+        return doc;
+    }
+
     @Override
     public SourceSection getSourceSection() {
         return body.getSourceSection();

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/nodes/NodeFactory.java

@@ -132,9 +132,9 @@ public <T> T duplicate(Node orig, Class<T> clazz) {
         return (T) NodeUtil.cloneNode(orig);
     }
 
-    public ModuleRootNode createModuleRoot(String name, ExpressionNode file, FrameDescriptor fd) {
+    public ModuleRootNode createModuleRoot(String name, String doc, ExpressionNode file, FrameDescriptor fd) {
         file.markAsRoot();
-        return new ModuleRootNode(language, name, file, fd, null);
+        return new ModuleRootNode(language, name, doc, file, fd, null);
     }
 
     public FunctionRootNode createFunctionRoot(SourceSection sourceSection, String functionName, boolean isGenerator, FrameDescriptor frameDescriptor, ExpressionNode body,

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/parser/PythonTreeTranslator.java

@@ -113,6 +113,8 @@ public final class PythonTreeTranslator extends Python3BaseVisitor<Object> {
     protected final AssignmentTranslator assigns;
     protected final Source source;
     protected final String name;
+    protected String moduleDoc;
+    protected boolean firstStatement = true;
 
     protected final boolean isInlineMode;
 
@@ -232,7 +234,7 @@ public Object visitFile_input(Python3Parser.File_inputContext ctx) {
         ExpressionNode file = asExpression(super.visitFile_input(ctx));
         deriveSourceSection(ctx, file);
         environment.popScope();
-        return factory.createModuleRoot(name, file, ctx.scope.getFrameDescriptor());
+        return factory.createModuleRoot(name, moduleDoc, file, ctx.scope.getFrameDescriptor());
     }
 
     @Override
@@ -256,7 +258,7 @@ public Object visitSingle_input(Python3Parser.Single_inputContext ctx) {
         if (isInlineMode) {
             return body;
         } else {
-            return factory.createModuleRoot("<expression>", body, ctx.scope.getFrameDescriptor());
+            return factory.createModuleRoot("<expression>", moduleDoc, body, ctx.scope.getFrameDescriptor());
         }
     }
 
@@ -1284,7 +1286,14 @@ private LoopNode createForInScope(StatementNode target, ExpressionNode iterator,
     @Override
     public Object visitExpr_stmt(Python3Parser.Expr_stmtContext ctx) {
         if (ctx.normassign().isEmpty() && ctx.annassign() == null && ctx.augassign() == null) {
-            return super.visitExpr_stmt(ctx);
+            Object exprNode = super.visitExpr_stmt(ctx);
+            if (firstStatement) {
+                firstStatement = false;
+                if (exprNode instanceof StringLiteralNode) {
+                    moduleDoc = ((StringLiteralNode) exprNode).getValue();
+                }
+            }
+            return exprNode;
         } else {
             return assigns.translate(ctx);
         }