Update src/ast/Python3VisitorImpl.java

suggested by santo

Co-authored-by: Santo Cariotti <santo@dcariotti.me>

1 files changed, 3 insertions, 4 deletions

diff --git a/src/ast/Python3VisitorImpl.java b/src/ast/Python3VisitorImpl.java
index bef75cf..42e2537 100644
--- a/src/ast/Python3VisitorImpl.java
+++ b/src/ast/Python3VisitorImpl.java
@@ -569,11 +569,10 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
         // TODO: Implement comp_if
         // Node iter = visit(ctx.comp_if());
         Comp_forContext cfc = ctx.comp_for();
+        Node forNode = null;
         if (cfc != null) {
-            Node forNode = visit(ctx.comp_for());
-            return new CompIterNode(forNode);
-
+            forNode = visit(ctx.comp_for());
         }
-        return new CompIterNode(null);
+        return new CompIterNode(forNode);
     }
 }