Pretty-printing nonlocal now works.

fabioz · Oct 1, 2023 · ab66dee · ab66dee
1 parent 1a1f7b7
commit ab66dee
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 8 deletions.
diff --git a/...ser/src/org/python/pydev/parser/prettyprinterv2/MakeAstValidForPrettyPrintingVisitor.java b/...ser/src/org/python/pydev/parser/prettyprinterv2/MakeAstValidForPrettyPrintingVisitor.java
@@ -54,6 +54,7 @@
 import org.python.pydev.parser.jython.ast.Module;
 import org.python.pydev.parser.jython.ast.Name;
 import org.python.pydev.parser.jython.ast.NameTok;
+import org.python.pydev.parser.jython.ast.NonLocal;
 import org.python.pydev.parser.jython.ast.Num;
 import org.python.pydev.parser.jython.ast.Pass;
 import org.python.pydev.parser.jython.ast.Print;
@@ -859,6 +860,14 @@ public Object visitClassDef(ClassDef node) throws Exception {
         return null;
     }
 
+    @Override
+    public Object visitNonLocal(NonLocal node) throws Exception {
+        fixNode(node);
+        node.traverse(this);
+        fixAfterNode(node);
+        return null;
+    }
+
     public boolean isFilled(SimpleNode[] nodes) {
         return (nodes != null) && (nodes.length > 0);
     }

diff --git a/...thon.pydev.parser/src/org/python/pydev/parser/prettyprinterv2/PrettyPrinterVisitorV2.java b/...thon.pydev.parser/src/org/python/pydev/parser/prettyprinterv2/PrettyPrinterVisitorV2.java
@@ -51,6 +51,7 @@
 import org.python.pydev.parser.jython.ast.MatchValue;
 import org.python.pydev.parser.jython.ast.Name;
 import org.python.pydev.parser.jython.ast.NameTok;
+import org.python.pydev.parser.jython.ast.NonLocal;
 import org.python.pydev.parser.jython.ast.Num;
 import org.python.pydev.parser.jython.ast.Pass;
 import org.python.pydev.parser.jython.ast.Print;
@@ -878,20 +879,54 @@ public Object visitStrJoin(StrJoin node) throws Exception {
     }
 
     @Override
-    public Object visitGlobal(Global node) throws Exception {
+    public Object visitNonLocal(NonLocal node) throws Exception {
         beforeNode(node);
-        doc.addRequire("global", node);
-
+        int id = doc.pushRecordChanges();
+        doc.addRequire("nonlocal ", node);
         if (node.names != null) {
-            for (int i = 0; i < node.names.length; i++) {
-                if (i > 0) {
-                    doc.addRequire(",", lastNode);
+            if (node.names.length > 0) {
+                for (int i = 0; i < node.names.length; i++) {
+                    if (i > 0) {
+                        doc.addRequire(",", lastNode);
+                    }
+                    if (node.names[i] != null) {
+                        node.names[i].accept(this);
+                    }
                 }
-                if (node.names[i] != null) {
-                    node.names[i].accept(this);
+            }
+        }
+        java.util.List<ILinePart> recordedChanges = this.doc.popRecordChanges(id);
+        this.doc.replaceRecorded(recordedChanges, "nonlocal", "nonlocal ");
+
+        if (node.value != null) {
+            doc.addRequire("=", lastNode);
+            node.value.accept(this);
+        }
+
+        afterNode(node);
+        return null;
+    }
+
+    @Override
+    public Object visitGlobal(Global node) throws Exception {
+        beforeNode(node);
+        int id = doc.pushRecordChanges();
+        doc.addRequire("global ", node);
+        if (node.names != null) {
+            if (node.names.length > 0) {
+                for (int i = 0; i < node.names.length; i++) {
+                    if (i > 0) {
+                        doc.addRequire(",", lastNode);
+                    }
+                    if (node.names[i] != null) {
+                        node.names[i].accept(this);
+                    }
                 }
             }
         }
+        java.util.List<ILinePart> recordedChanges = this.doc.popRecordChanges(id);
+        this.doc.replaceRecorded(recordedChanges, "global", "global ");
+
         if (node.value != null) {
             doc.addRequire("=", lastNode);
             node.value.accept(this);

diff --git a/....python.pydev.parser/tests/org/python/pydev/parser/prettyprinter/PrettyPrinter30Test.java b/....python.pydev.parser/tests/org/python/pydev/parser/prettyprinter/PrettyPrinter30Test.java
@@ -413,6 +413,25 @@ public void testOthers1() throws Exception {
         checkPrettyPrintEqual(s, expected);
     }
 
+    public void testNonlocal() throws Exception {
+        String s = ""
+                + "def another():\n"
+                + "    a,b,c = 3\n"
+                + "    def method():\n"
+                + "        nonlocal a\n"
+                + "        nonlocal b,c\n";
+
+        checkPrettyPrintEqual(s);
+    }
+
+    public void testGlobal() throws Exception {
+        String s = "def method():\n"
+                + "    global a\n"
+                + "    global b,c\n";
+
+        checkPrettyPrintEqual(s);
+    }
+
     public void testOthers2() throws Throwable {
         final String s = "" +
                 "def _format(node):\n" +