Move mermaid rendering into its own file

Author: kddnewton

lib/syntax_tree.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-require "cgi"
 require "etc"
 require "json"
 require "pp"
@@ -71,19 +70,6 @@ module SyntaxTree
   # that Syntax Tree can format arbitrary parts of a document.
   DEFAULT_INDENTATION = 0
 
-  # This is a hook provided so that plugins can register themselves as the
-  # handler for a particular file type.
-  def self.register_handler(extension, handler)
-    HANDLERS[extension] = handler
-  end
-
-  # Parses the given source and returns the syntax tree.
-  def self.parse(source)
-    parser = Parser.new(source)
-    response = parser.parse
-    response unless parser.error?
-  end
-
   # Parses the given source and returns the formatted source.
   def self.format(
     source,
@@ -98,13 +84,39 @@ def self.format(
     formatter.output.join
   end
 
+  # Indexes the given source code to return a list of all class, module, and
+  # method definitions. Used to quickly provide indexing capability for IDEs or
+  # documentation generation.
+  def self.index(source)
+    Index.index(source)
+  end
+
+  # Indexes the given file to return a list of all class, module, and method
+  # definitions. Used to quickly provide indexing capability for IDEs or
+  # documentation generation.
+  def self.index_file(filepath)
+    Index.index_file(filepath)
+  end
+
   # A convenience method for creating a new mutation visitor.
   def self.mutation
     visitor = Visitor::MutationVisitor.new
     yield visitor
     visitor
   end
 
+  # Parses the given source and returns the syntax tree.
+  def self.parse(source)
+    parser = Parser.new(source)
+    response = parser.parse
+    response unless parser.error?
+  end
+
+  # Parses the given file and returns the syntax tree.
+  def self.parse_file(filepath)
+    parse(read(filepath))
+  end
+
   # Returns the source from the given filepath taking into account any potential
   # magic encoding comments.
   def self.read(filepath)
@@ -120,23 +132,15 @@ def self.read(filepath)
     File.read(filepath, encoding: encoding)
   end
 
+  # This is a hook provided so that plugins can register themselves as the
+  # handler for a particular file type.
+  def self.register_handler(extension, handler)
+    HANDLERS[extension] = handler
+  end
+
   # Searches through the given source using the given pattern and yields each
   # node in the tree that matches the pattern to the given block.
   def self.search(source, query, &block)
     Search.new(Pattern.new(query).compile).scan(parse(source), &block)
   end
-
-  # Indexes the given source code to return a list of all class, module, and
-  # method definitions. Used to quickly provide indexing capability for IDEs or
-  # documentation generation.
-  def self.index(source)
-    Index.index(source)
-  end
-
-  # Indexes the given file to return a list of all class, module, and method
-  # definitions. Used to quickly provide indexing capability for IDEs or
-  # documentation generation.
-  def self.index_file(filepath)
-    Index.index_file(filepath)
-  end
 end

lib/syntax_tree/mermaid.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+require "cgi"
+
+module SyntaxTree
+  # This module is responsible for rendering mermaid flow charts.
+  module Mermaid
+    class Node
+      SHAPES = %i[circle rectangle stadium].freeze
+
+      attr_reader :id, :label, :shape
+
+      def initialize(id, label, shape)
+        raise unless SHAPES.include?(shape)
+
+        @id = id
+        @label = label
+        @shape = shape
+      end
+
+      def render
+        left_bound, right_bound =
+          case shape
+          when :circle
+            ["((", "))"]
+          when :rectangle
+            ["[", "]"]
+          when :stadium
+            ["([", "])"]
+          end
+
+        "  #{id}#{left_bound}\"#{CGI.escapeHTML(label)}\"#{right_bound}"
+      end
+    end
+
+    class Edge
+      TYPES = %i[directed].freeze
+
+      attr_reader :from, :to, :label, :type
+
+      def initialize(from, to, label, type)
+        raise unless TYPES.include?(type)
+
+        @from = from
+        @to = to
+        @label = label
+        @type = type
+      end
+
+      def render
+        case type
+        when :directed
+          "  #{from.id} -- \"#{CGI.escapeHTML(label)}\" --> #{to.id}"
+        end
+      end
+    end
+
+    class FlowChart
+      attr_reader :nodes, :edges
+
+      def initialize
+        @nodes = {}
+        @edges = []
+      end
+
+      def edge(from, to, label, type = :directed)
+        edges << Edge.new(from, to, label, type)
+      end
+
+      def node(id, label, shape = :rectangle)
+        nodes[id] = Node.new(id, label, shape)
+      end
+
+      def render
+        output = StringIO.new
+        output.puts("flowchart TD")
+
+        nodes.each_value { |node| output.puts(node.render) }
+        edges.each { |edge| output.puts(edge.render) }
+
+        output.string
+      end
+    end
+  end
+end

lib/syntax_tree/visitor/mermaid_visitor.rb

@@ -4,18 +4,16 @@ module SyntaxTree
   class Visitor
     # This visitor transforms the AST into a mermaid flow chart.
     class MermaidVisitor < FieldVisitor
-      attr_reader :output, :target
+      attr_reader :flowchart, :target
 
       def initialize
-        @output = StringIO.new
-        @output.puts("flowchart TD")
-
+        @flowchart = Mermaid::FlowChart.new
         @target = nil
       end
 
       def visit_program(node)
         super
-        output.string
+        flowchart.render
       end
 
       private
@@ -26,19 +24,13 @@ def comments(node)
 
       def field(name, value)
         case value
-        when Node
-          node_id = visit(value)
-          output.puts("  #{target} -- \"#{name}\" --> #{node_id}")
-        when String
-          node_id = "#{target}_#{name}"
-          output.puts("  #{node_id}([#{CGI.escapeHTML(value.inspect)}])")
-          output.puts("  #{target} -- \"#{name}\" --> #{node_id}")
         when nil
           # skip
+        when Node
+          flowchart.edge(target, visit(value), name)
         else
-          node_id = "#{target}_#{name}"
-          output.puts("  #{node_id}([\"#{CGI.escapeHTML(value.inspect)}\"])")
-          output.puts("  #{target} -- \"#{name}\" --> #{node_id}")
+          to = flowchart.node("#{target.id}_#{name}", value.inspect, :stadium)
+          flowchart.edge(target, to, name)
         end
       end
 
@@ -52,11 +44,8 @@ def node(node, type)
         previous_target = target
 
         begin
-          @target = "node_#{node.object_id}"
-
+          @target = flowchart.node("node_#{node.object_id}", type)
           yield
-
-          output.puts("  #{@target}[\"#{type}\"]")
           @target
         ensure
           @target = previous_target
@@ -65,11 +54,11 @@ def node(node, type)
 
       def pairs(name, values)
         values.each_with_index do |(key, value), index|
-          node_id = "#{target}_#{name}_#{index}"
-          output.puts("  #{node_id}((\"&nbsp;\"))")
-          output.puts("  #{target} -- \"#{name}[#{index}]\" --> #{node_id}")
-          output.puts("  #{node_id} -- \"[0]\" --> #{visit(key)}")
-          output.puts("  #{node_id} -- \"[1]\" --> #{visit(value)}") if value
+          to = flowchart.node("#{target.id}_#{name}_#{index}", " ", :circle)
+
+          flowchart.edge(target, to, "#{name}[#{index}]")
+          flowchart.edge(to, visit(key), "[0]")
+          flowchart.edge(to, visit(value), "[1]") if value
         end
       end