SigDivider divides a single RBI file into multiple files by each class or module.

sig_divider.rb

# frozen_string_literal: true

require 'fileutils'
require 'rbi'
require 'active_support/all'

# SigDivider divides a single RBI file into multiple files by each class or module.
class SigDivider < RBI::Visitor
  TreeWithName = Struct.new(:name, :tree, keyword_init: true)

  # @return [String]
  attr_reader :output_dir

  # @param output_dir [String]
  def initialize(output_dir:)
    @output_dir = output_dir
  end

  # @param rbi_source [String]
  def process(rbi_source:)
    rbi_tree = RBI::Parser.parse_string(rbi_source)
    trees_with_names = visit(rbi_tree)

    if rbi_tree.nodes.present?
      trees_with_names << TreeWithName.new(name: "Root", tree: rbi_tree)
    end

    FileUtils.rmtree(output_dir)

    trees_with_names.each do |tree_with_name|
      filename = File.join(output_dir, "#{tree_with_name.name.underscore}.rbi")

      RBI::Formatter.new(
        group_nodes: true,
        sort_nodes: true,
      ).format_tree(tree_with_name.tree)

      puts "Writing #{tree_with_name.name} -> #{filename}"
      FileUtils.mkdir_p(File.dirname(filename))

      tree_string = <<~RBI
      # typed: true

      # DO NOT EDIT MANUALLY.
      # THIS FILE IS AUTOGENERATED.

      #{tree_with_name.tree.string.chomp}
      RBI

      File.write(filename, tree_string)
    end
  end

  # @param node [RBI::Node, nil]
  # @return [Array<TreeWithName>]
  def visit(node)
    case node
    when RBI::Module
      # puts "Module: #{node.fully_qualified_name}"
      tree = build_detached_tree(node)
      [tree, *node.nodes.dup.flat_map { |n| visit(n) }]
    when RBI::Class
      # puts "Class: #{node.fully_qualified_name}"
      tree = build_detached_tree(node)
      [tree, *node.nodes.dup.flat_map { |n| visit(n) }]
    when RBI::Tree
      # puts "Tree: #{node.try(:index_ids)}"
      node.nodes.dup.flat_map { |n| visit(n) }
    else
      # puts "Other: #{node.try(:index_ids)}"
      []
    end
  end

  # @param node [RBI::Scope, nil]
  # @return [TrewWithName]
  def build_detached_tree(node)
    name = node.fully_qualified_name
    ancestors_rev = node_ancestors(node).reverse
    node.detach

    tree = ancestors_rev.reduce(node) do |child_node, parent_node_origin|
      parent_node =
        case parent_node_origin
        when RBI::Module
          RBI::Module.new(parent_node_origin.name)
        when RBI::Class
          RBI::Class.new(parent_node_origin.name)
        when RBI::SingletonClass
          RBI::SingletonClass.new
        when RBI::Struct
          RBI::Struct.new(parent_node_origin.name, members: parent_node_origin.members, keyword_init: parent_node_origin.keyword_init)
        end

      parent_node << child_node
      parent_node
    end

    TreeWithName.new(name: name, tree: tree)
  end

  # @param node [RBI::Scope, nil]
  # @return [Array<RBI::Scope>]
  def node_ancestors(node)
    ancestors = []
    iter = node
    while iter = iter.parent_scope
      ancestors.unshift(iter)
    end
    ancestors
  end
end

SigDivider.new(output_dir: ARGV[0]).process(rbi_source: File.read(ARGV[1]))