Separate statements by separate_statements (int) newlines.

pglast/printer.py

@@ -180,9 +180,8 @@ class RawStream(OutputStream):
 
     :param int expression_level:
            start the stream with the given expression level depth, 0 by default
-    :param bool separate_statements:
-           ``True`` by default, tells whether multiple statements shall be separated by an
-           empty line
+    :param int separate_statements:
+           ``1`` by default, tells how many empty lines should separate statements
     :param bool special_functions:
            ``False`` by default, when ``True`` some functions are treated in a special way and
            emitted as equivalent constructs
@@ -195,7 +194,7 @@ class RawStream(OutputStream):
     without any adornment.
     """
 
-    def __init__(self, expression_level=0, separate_statements=True, special_functions=False,
+    def __init__(self, expression_level=0, separate_statements=1, special_functions=False,
                  comma_at_eoln=False):
         super().__init__()
         self.expression_level = expression_level
@@ -227,7 +226,7 @@ def __call__(self, sql, plpgsql=False):
             else:
                 self.write(';')
                 self.newline()
-                if self.separate_statements:
+                for _ in range(int(self.separate_statements)):
                     self.newline()
             self.print_node(statement)
         return self.getvalue()

tests/test_printer.py

@@ -271,6 +271,24 @@ def raw_stmt(node, output):
             printer.NODE_PRINTERS.pop('RawStmt', None)
 
 
+def test_separate_statements():
+    """Separate statements by ``separate_statements`` (int) newlines."""
+    raw_stmt_printer = printer.NODE_PRINTERS.pop('RawStmt', None)
+    try:
+        @printer.node_printer('RawStmt')
+        def raw_stmt(node, output):
+            output.write('Yeah')
+
+        output = printer.IndentedStream(separate_statements=2)
+        result = output('SELECT 1; SELECT 2')
+        assert result == 'Yeah;\n\n\nYeah'
+    finally:
+        if raw_stmt_printer is not None:
+            printer.NODE_PRINTERS['RawStmt'] = raw_stmt_printer
+        else:
+            printer.NODE_PRINTERS.pop('RawStmt', None)
+
+
 def test_special_function():
     output = printer.RawStream(special_functions=True)