fix #24: fix block indent when rendering lua

boolangery · May 24, 2022 · 0b3afea · 0b3afea
1 parent 88a3ada
commit 0b3afea
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 9 deletions.
diff --git a/luaparser/printers.py b/luaparser/printers.py
@@ -218,6 +218,7 @@ def visit(self, node):
 class LuaOutputVisitor:
     def __init__(self, indent_size: int):
         self._indent_size = indent_size
+        self._level = 0
 
     @visitor(str)
     def visit(self, node) -> str:
@@ -245,9 +246,11 @@ def visit(self, node) -> str:
 
     @visitor(Block)
     def visit(self, node: Block) -> str:
+        self._level += 1
         output = indent(
-            "\n".join([self.visit(n) for n in node.body]), " " * self._indent_size
+            "\n".join([self.visit(n) for n in node.body]), " " * (self._indent_size if self._level > 1 else 0)
         )
+        self._level -= 1
         return output
 
     @visitor(Assign)
@@ -274,21 +277,21 @@ def visit(self, node: If) -> str:
             "if " + self.visit(node.test) + " then\n" + self.visit(node.body)
         )
         if isinstance(node.orelse, ElseIf):
-            output += self.visit(node.orelse)
+            output += "\n" + self.visit(node.orelse)
         elif node.orelse:
-            output += "else\n" + self.visit(node.orelse)
+            output += "\nelse\n" + self.visit(node.orelse)
         output += "\nend"
         return output
 
     @visitor(ElseIf)
     def visit(self, node: ElseIf) -> str:
         output = (
-            "elseif " + self.visit(node.test) + " then\n" + self.visit(node.body) + "\n"
+            "elseif " + self.visit(node.test) + " then\n" + self.visit(node.body)
         )
         if isinstance(node.orelse, ElseIf):
-            output += self.visit(node.orelse)
-        else:
-            output += "else\n" + self.visit(node.orelse)
+            output += "\n" + self.visit(node.orelse)
+        elif node.orelse:
+            output += "\nelse\n" + self.visit(node.orelse)
         return output
 
     @visitor(Label)

diff --git a/luaparser/tests/test_lua_output.py b/luaparser/tests/test_lua_output.py
@@ -129,10 +129,10 @@ def test_int_24(self):
                 if a == 1 then
                     if a == 2 then
                         if a == 3 then
+                        
                         end
                     end
                 end
-            end
-            """
+            end"""
         )
         self.assertEqual(source, ast.to_lua_source(ast.parse(source)))