From ad92929a1025cf38eeaee88138f132f36bec2291 Mon Sep 17 00:00:00 2001
From: Roberto Alsina <roberto.alsina@gmail.com>
Date: Mon, 19 Aug 2024 09:59:01 -0300
Subject: [PATCH] Make formatter a bit more convenient

---
 src/formatter.cr       | 2 +-
 src/formatters/ansi.cr | 4 +++-
 src/formatters/html.cr | 4 +++-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/formatter.cr b/src/formatter.cr
index 12e0814..7c0bf44 100644
--- a/src/formatter.cr
+++ b/src/formatter.cr
@@ -12,7 +12,7 @@ module Tartrazine
     property theme : Theme = Tartrazine.theme("default-dark")
 
     # Format the text using the given lexer.
-    def format(text : String, lexer : Lexer) : String
+    def format(text : String, lexer : Lexer, io : IO?) : String?
       raise Exception.new("Not implemented")
     end
 
diff --git a/src/formatters/ansi.cr b/src/formatters/ansi.cr
index c1bea9a..fb3e5c5 100644
--- a/src/formatters/ansi.cr
+++ b/src/formatters/ansi.cr
@@ -11,7 +11,8 @@ module Tartrazine
       "#{i + 1}".rjust(4).ljust(5)
     end
 
-    def format(text : String, lexer : Lexer, outp : IO) : Nil
+    def format(text : String, lexer : Lexer, io : IO?) : String?
+      outp = io.nil? ? String::Builder.new("") : io
       tokenizer = Tokenizer.new(lexer, text)
       i = 0
       outp << line_label(i) if line_numbers?
@@ -22,6 +23,7 @@ module Tartrazine
           outp << line_label(i) if line_numbers?
         end
       end
+      return outp.to_s if io.nil?
     end
 
     def colorize(text : String, token : String) : String
diff --git a/src/formatters/html.cr b/src/formatters/html.cr
index 3b994d1..ac9aa25 100644
--- a/src/formatters/html.cr
+++ b/src/formatters/html.cr
@@ -58,7 +58,8 @@ module Tartrazine
       "<span #{line_id} #{line_class} style=\"user-select: none;\">#{line_label} </span>"
     end
 
-    def format_text(text : String, lexer : Lexer, outp : IO) : Nil
+    def format_text(text : String, lexer : Lexer, io : IO?) : String?
+      outp = io.nil? ? String::Builder.new("") : io
       tokenizer = Tokenizer.new(lexer, text)
       i = 0
       if surrounding_pre?
@@ -75,6 +76,7 @@ module Tartrazine
         end
       end
       outp << "</code></pre>"
+      return outp.to_s if io.nil?
     end
 
     # ameba:disable Metrics/CyclomaticComplexity