first cut rbi generator

Author: drewhoffer

lib/rbi_generator/generator.rb

@@ -13,29 +13,79 @@ def initialize(file_path)
 
     def generate
       content = File.read(@file_path)
-      ast = Ripper.sexp(content)
-      pp ast
-      class_name = extract_class_name(ast)
-      method_names = extract_method_names(ast)
-      puts "Class Name: #{class_name}"
-      puts "Method Names: #{method_names.join(", ")}"
-      # puts "Method Names: #{method_names.join(", ")}"
-      # puts Ripper.sexp(content)
+      result = Prism.parse(content)
+      class_node = find_first_class_node(result.value)
+      return unless class_node
+
+      class_name = class_node.name
+      method_names = get_all_method_names(class_node)
+
+      write_rbi_to_file(class_name, method_names)
+      puts "Generated RBI file: #{@rbi_path}"
     end
 
-    # Takes the syntax tree and extracts the class name
-    def extract_class_name(ast)
-      class_node = ast[1].find { |node| node[0] == :class }
-      class_node[1][1] if class_node
+    private
+
+    def write_rbi_to_file(class_name, methods)
+      rbi = +"class #{class_name}\n"
+
+      methods.each do |method|
+        name = method[:name]
+        params = method[:params]
+        sig = if params.empty?
+                "  sig { void }"
+              else
+                typed_params = params.map { |p| "#{p}: T.untyped" }.join(", ")
+                "  sig { params(#{typed_params}).returns(T.untyped) }"
+              end
+        rbi << "#{sig}\n"
+        def_line = "  def #{name}"
+        def_line << "(#{params.join(", ")})" unless params.empty?
+        def_line << "; end\n\n"
+        rbi << def_line
+      end
+
+      rbi << "end\n"
+      File.write(@rbi_path, rbi)
     end
 
-    def extract_method_names(ast)
-      class_node = ast[1].find { |node| node[0] == :class }
-      body = class_node[3] # :bodystmt
-      return [] unless body && body[0].is_a?(Array)
+    # Recursively find the first ClassNode in the AST
+    def find_first_class_node(node)
+      return node if node.is_a?(Prism::ClassNode)
+      return nil unless node.respond_to?(:child_nodes)
+
+      node.child_nodes.each do |child|
+        found = find_first_class_node(child)
+        return found if found
+      end
+    end
+
+    # Goes through the class and grab every child node (they will be defined as defNode)
+    def get_all_method_names(class_node)
+      class_node.body.child_nodes
+                .select { |child| child.is_a?(Prism::DefNode) }
+                .map do |method_node|
+        {
+          name: method_node.name,
+          params: get_method_params(method_node)
+        }
+      end
+    end
+
+    def get_method_params(method_node)
+      params_node = method_node.parameters
+      return [] unless params_node
+
+      all_params = []
 
-      method_defs = body[0].select { |node| node[0] == :def }
-      method_defs.map { |def_node| def_node[1][1] } # => ["hello", "goodbye"]
+      all_params += params_node.requireds.map { |p| p.name }
+      all_params += params_node.optionals.map { |p| p.name }
+      all_params << params_node.rest.name if params_node.rest
+      all_params += params_node.posts.map { |p| p.name }
+      all_params += params_node.keywords.map { |p| p.name }
+      all_params << params_node.keyword_rest.name if params_node.keyword_rest
+      all_params << params_node.block.name if params_node.block
+      all_params.compact
     end
   end
 end

test.rb

@@ -2,4 +2,6 @@ class MyTest
   def hello; end
 
   def test(x); end
+
+  def add(x, y); end
 end

test.rbi

@@ -0,0 +1,10 @@
+class MyTest
+  sig { void }
+  def hello; end
+
+  sig { params(x: T.untyped).returns(T.untyped) }
+  def test(x); end
+
+  sig { params(x: T.untyped, y: T.untyped).returns(T.untyped) }
+  def add(x, y); end
+end