Format non-Ruby STDIN/script content too

lib/syntax_tree/cli.rb

@@ -63,12 +63,13 @@ def writable?
     class ScriptItem
       attr_reader :source
 
-      def initialize(source)
+      def initialize(source, handler_extension)
         @source = source
+        @handler_extension = handler_extension
       end
 
       def handler
-        HANDLERS[".rb"]
+        HANDLERS[@handler_extension]
       end
 
       def filepath
@@ -82,8 +83,12 @@ def writable?
 
     # An item of work that correspond to the content passed in via stdin.
     class STDINItem
+      def initialize(handler_extension)
+        @handler_extension = handler_extension
+      end
+
       def handler
-        HANDLERS[".rb"]
+        HANDLERS[@handler_extension]
       end
 
       def filepath
@@ -457,7 +462,10 @@ def run(item)
         The maximum line width to use when formatting.
 
       -e SCRIPT
-        Parse an inline Ruby string.
+        Parse an inline string.
+
+      --handler=EXTENSION
+        A file extension matching the content passed in via STDIN or -e. Defaults to 'rb'
     HELP
 
     # This represents all of the options that can be passed to the CLI. It is
@@ -468,13 +476,15 @@ class Options
                   :plugins,
                   :print_width,
                   :scripts,
+                  :handler_extension,
                   :target_ruby_version
 
       def initialize
         @ignore_files = []
         @plugins = []
         @print_width = DEFAULT_PRINT_WIDTH
         @scripts = []
+        @handler_extension = ".rb"
         @target_ruby_version = DEFAULT_RUBY_VERSION
       end
 
@@ -523,6 +533,13 @@ def parser
           # it and add it to the list of scripts to run.
           opts.on("-e SCRIPT") { |script| @scripts << script }
 
+          # If there is a handler specified, then parse it and use it for
+          # STDIN and scripts.
+          opts.on("--handler=EXTENSION") do |extension|
+            # Both ".rb" and "rb" are going to work
+            @handler_extension = ".#{extension.delete_prefix(".")}"
+          end
+
           # If there is a target ruby version specified on the command line,
           # parse that out and use it when formatting.
           opts.on("--target-ruby-version=VERSION") do |version|
@@ -630,9 +647,11 @@ def run(argv)
               end
           end
 
-          options.scripts.each { |script| queue << ScriptItem.new(script) }
+          options.scripts.each do |script|
+            queue << ScriptItem.new(script, options.handler_extension)
+          end
         else
-          queue << STDINItem.new
+          queue << STDINItem.new(options.handler_extension)
         end
 
         # At the end, we're going to return whether or not this worker ever

test/cli_test.rb

@@ -10,6 +10,10 @@ def parse(source)
         source * 2
       end
 
+      def format(source, _print_width, **)
+        "Formatted #{source}"
+      end
+
       def read(filepath)
         File.read(filepath)
       end
@@ -202,6 +206,26 @@ def test_multiple_inline_scripts
       assert_equal(["1 + 1", "2 + 2"], stdio.split("\n").sort)
     end
 
+    def test_format_script_with_custom_handler
+      SyntaxTree.register_handler(".test", TestHandler.new)
+      stdio, =
+        capture_io { SyntaxTree::CLI.run(%w[format --handler=test -e <test>]) }
+      assert_equal("Formatted <test>\n", stdio)
+    ensure
+      SyntaxTree::HANDLERS.delete(".test")
+    end
+
+    def test_format_stdin_with_custom_handler
+      SyntaxTree.register_handler(".test", TestHandler.new)
+      stdin = $stdin
+      $stdin = StringIO.new("<test>")
+      stdio, = capture_io { SyntaxTree::CLI.run(%w[format --handler=test]) }
+      assert_equal("Formatted <test>\n", stdio)
+    ensure
+      $stdin = stdin
+      SyntaxTree::HANDLERS.delete(".test")
+    end
+
     def test_generic_error
       SyntaxTree.stub(:format, ->(*) { raise }) do
         result = run_cli("format")