Mess with wrapping and whatnot

grammar.py

@@ -25,6 +25,8 @@ class FineGrammar(Grammar):
 
     trivia = ["BLANKS", "COMMENT"]
 
+    pretty_indent = "  "
+
     def __init__(self):
         super().__init__(
             precedence=[
@@ -134,14 +136,23 @@ class FineGrammar(Grammar):
     @rule("FunctionDecl")
     def function_declaration(self) -> Rule:
         return seq(
+            group(
+                group(
                     group(
                         self.FUN,
                         sp,
-                mark(self.IDENTIFIER, field="name", highlight=highlight.entity.name.function),
+                        mark(
-                sp,
+                            self.IDENTIFIER,
-                group(
+                            field="name",
+                            highlight=highlight.entity.name.function,
+                        ),
+                    ),
+                    nl,
                     mark(self.function_parameters, field="parameters"),
-                    mark(opt(sp, group(self.ARROW, sp, self.type_expression)), field="return_type"),
+                ),
+                mark(
+                    opt(indent(sp, group(self.ARROW, sp, self.type_expression))),
+                    field="return_type",
                 ),
             ),
             sp,

harness.py

@@ -580,7 +580,9 @@ class Harness:
             return []
 
         return (
-            wadler.layout_document(self.document, width).apply_to_source(self.source).splitlines()
+            wadler.layout_document(self.document, width, self.load_printer().indent())
+            .apply_to_source(self.source)
+            .splitlines()
         )
 
 

parser/wadler.py

@@ -108,7 +108,7 @@ class DocumentLayout:
         return result
 
 
-def layout_document(doc: Document, width: int) -> DocumentLayout:
+def layout_document(doc: Document, width: int, indent: str) -> DocumentLayout:
     """Lay out a document to fit within the given width.
 
     The result of this function is a DocumentLayout which can trivially be
@@ -213,13 +213,13 @@ def layout_document(doc: Document, width: int) -> DocumentLayout:
                     column += len(replace)
                 else:
                     # TODO: Custom newline expansion, custom indent segments.
-                    output.append("\n" + (chunk.indent * " "))
+                    output.append("\n" + (chunk.indent * indent))
-                    column = chunk.indent
+                    column = chunk.indent * len(indent)
 
             case ForceBreak():
                 # TODO: Custom newline expansion, custom indent segments.
-                output.append("\n" + (chunk.indent * " "))
+                output.append("\n" + (chunk.indent * indent))
-                column = chunk.indent
+                column = chunk.indent * len(indent)
 
             case Cons(left, right):
                 chunks.append(chunk.with_document(right))
@@ -363,12 +363,20 @@ class Printer:
     grammar: parser.Grammar
     _matchers: dict[str, Matcher]
     _nonterminals: dict[str, parser.NonTerminal]
+    _indent: str
 
-    def __init__(self, grammar: parser.Grammar):
+    def __init__(self, grammar: parser.Grammar, indent: str | None = None):
         self.grammar = grammar
         self._nonterminals = {nt.name: nt for nt in grammar.non_terminals()}
         self._matchers = {}
 
+        if indent is None:
+            indent = getattr(self.grammar, "pretty_indent", " ")
+        self._indent = indent
+
+    def indent(self) -> str:
+        return self._indent
+
     def lookup_nonterminal(self, name: str) -> parser.NonTerminal:
         return self._nonterminals[name]
 
@@ -552,4 +560,4 @@ class Printer:
 
     def format_tree(self, tree: runtime.Tree, width: int) -> DocumentLayout:
         doc = self.convert_tree_to_document(tree)
-        return layout_document(doc, width)
+        return layout_document(doc, width, self._indent)