[Tapioca Addon] Support gem RBI generation

lib/ruby_lsp/tapioca/addon.rb

@@ -13,6 +13,7 @@
 end
 
 require "zlib"
+require "ruby_lsp/tapioca/run_gem_rbi_check"
 
 module RubyLsp
   module Tapioca
@@ -27,6 +28,7 @@ def initialize
         @rails_runner_client = T.let(nil, T.nilable(RubyLsp::Rails::RunnerClient))
         @index = T.let(nil, T.nilable(RubyIndexer::Index))
         @file_checksums = T.let({}, T::Hash[String, String])
+        @lockfile_diff = T.let(nil, T.nilable(String))
         @outgoing_queue = T.let(nil, T.nilable(Thread::Queue))
       end
 
@@ -50,6 +52,8 @@ def activate(global_state, outgoing_queue)
             request_name: "load_compilers_and_extensions",
             workspace_path: @global_state.workspace_path,
           )
+
+          run_gem_rbi_check
         rescue IncompatibleApiError
           # The requested version for the Rails add-on no longer matches. We need to upgrade and fix the breaking
           # changes
@@ -132,6 +136,20 @@ def file_updated?(change, path)
 
         false
       end
+
+      sig { void }
+      def run_gem_rbi_check
+        gem_rbi_check = RunGemRbiCheck.new
+        gem_rbi_check.run
+
+        T.must(@outgoing_queue) << Notification.window_log_message(
+          gem_rbi_check.stdout,
+        ) unless gem_rbi_check.stdout.empty?
+        T.must(@outgoing_queue) << Notification.window_log_message(
+          gem_rbi_check.stderr,
+          type: Constant::MessageType::WARNING,
+        ) unless gem_rbi_check.stderr.empty?
+      end
     end
   end
 end

lib/ruby_lsp/tapioca/lockfile_diff_parser.rb

@@ -0,0 +1,49 @@
+# typed: true
+# frozen_string_literal: true
+
+require "bundler"
+
+module RubyLsp
+  module Tapioca
+    class LockfileDiffParser
+      GEM_NAME_PATTERN = /[\w\-]+/
+      DIFF_LINE_PATTERN = /[+-](.*#{GEM_NAME_PATTERN})\s*\(/
+      ADDED_LINE_PATTERN = /^\+.*#{GEM_NAME_PATTERN} \(.*\)/
+      REMOVED_LINE_PATTERN = /^-.*#{GEM_NAME_PATTERN} \(.*\)/
+
+      attr_reader :added_or_modified_gems
+      attr_reader :removed_gems
+
+      def initialize(diff_content, direct_dependencies: nil)
+        @diff_content = diff_content.lines
+        @current_dependencies = direct_dependencies ||
+          Bundler::LockfileParser.new(Bundler.default_lockfile.read).dependencies.keys
+        @added_or_modified_gems = parse_added_or_modified_gems
+        @removed_gems = parse_removed_gems
+      end
+
+      private
+
+      def parse_added_or_modified_gems
+        @diff_content
+          .filter_map { |line| extract_gem(line) if line.match?(ADDED_LINE_PATTERN) }
+          .uniq
+      end
+
+      def parse_removed_gems
+        @diff_content.filter_map do |line|
+          next unless line.match?(REMOVED_LINE_PATTERN)
+
+          gem = extract_gem(line)
+          next if @current_dependencies.include?(gem)
+
+          gem
+        end.uniq
+      end
+
+      def extract_gem(line)
+        line.match(DIFF_LINE_PATTERN)[1].strip
+      end
+    end
+  end
+end

lib/ruby_lsp/tapioca/run_gem_rbi_check.rb

@@ -0,0 +1,116 @@
+# typed: true
+# frozen_string_literal: true
+
+require "open3"
+require "ruby_lsp/tapioca/lockfile_diff_parser"
+
+module RubyLsp
+  module Tapioca
+    class RunGemRbiCheck
+      extend T::Sig
+
+      attr_reader :stdout
+      attr_reader :stderr
+      attr_reader :status
+
+      sig { void }
+      def initialize
+        @stdout = T.let("", String)
+        @stderr = T.let("", String)
+        @status = T.let(nil, T.nilable(Process::Status))
+      end
+
+      sig { params(project_path: String).void }
+      def run(project_path = ".")
+        FileUtils.chdir(project_path) do
+          return log_message("Not a git repository") unless git_repo?
+
+          lockfile_changed? ? generate_gem_rbis : cleanup_orphaned_rbis
+        end
+      end
+
+      private
+
+      sig { returns(T::Boolean) }
+      def git_repo?
+        _, status = Open3.capture2e("git rev-parse --is-inside-work-tree")
+        T.must(status.success?)
+      end
+
+      sig { returns(T::Boolean) }
+      def lockfile_changed?
+        fetch_lockfile_diff
+        !@lockfile_diff.empty?
+      end
+
+      sig { returns(String) }
+      def fetch_lockfile_diff
+        @lockfile_diff = File.exist?("Gemfile.lock") ? %x(git diff Gemfile.lock).strip : ""
+      end
+
+      sig { void }
+      def generate_gem_rbis
+        parser = Tapioca::LockfileDiffParser.new(@lockfile_diff)
+        removed_gems = parser.removed_gems
+        added_or_modified_gems = parser.added_or_modified_gems
+
+        if added_or_modified_gems.any?
+          log_message("Identified lockfile changes, attempting to generate gem RBIs...")
+          execute_tapioca_gem_command(added_or_modified_gems)
+        elsif removed_gems.any?
+          remove_rbis(removed_gems)
+        end
+      end
+
+      sig { params(gems: T::Array[String]).void }
+      def execute_tapioca_gem_command(gems)
+        Bundler.with_unbundled_env do
+          stdout, stderr, status = T.unsafe(Open3).capture3(
+            "bundle",
+            "exec",
+            "tapioca",
+            "gem",
+            "--lsp_addon",
+            *gems,
+          )
+
+          log_message(stdout) unless stdout.empty?
+          @stderr = stderr unless stderr.empty?
+          @status = status
+        end
+      end
+
+      sig { params(gems: T::Array[String]).void }
+      def remove_rbis(gems)
+        FileUtils.rm_f(Dir.glob("sorbet/rbi/gems/{#{gems.join(",")}}@*.rbi"))
+        log_message("Removed RBIs for: #{gems.join(", ")}")
+      end
+
+      sig { void }
+      def cleanup_orphaned_rbis
+        untracked_files = %x(git ls-files --others --exclude-standard sorbet/rbi/gems/).lines.map(&:strip)
+        deleted_files = %x(git ls-files --deleted sorbet/rbi/gems/).lines.map(&:strip)
+
+        delete_files(untracked_files, "Deleted untracked RBIs")
+        restore_files(deleted_files, "Restored deleted RBIs")
+      end
+
+      sig { params(files: T::Array[String], message: String).void }
+      def delete_files(files, message)
+        files.each { |file| File.delete(file) }
+        log_message("#{message}: #{files.join(", ")}") unless files.empty?
+      end
+
+      sig { params(files: T::Array[String], message: String).void }
+      def restore_files(files, message)
+        files.each { |file| %x(git checkout -- #{file}) }
+        log_message("#{message}: #{files.join(", ")}") unless files.empty?
+      end
+
+      sig { params(message: String).void }
+      def log_message(message)
+        @stdout += "#{message}\n"
+      end
+    end
+  end
+end

spec/helpers/mock_gem.rb

@@ -66,5 +66,11 @@ def default_gemspec_contents
         end
       GEMSPEC
     end
+
+    sig { params(version: String).void }
+    def update(version)
+      @version = version
+      gemspec(default_gemspec_contents)
+    end
   end
 end

spec/tapioca/ruby_lsp/run_gem_rbi_check_spec.rb

@@ -0,0 +1,141 @@
+# typed: true
+# frozen_string_literal: true
+
+require "spec_helper"
+require "ruby_lsp/tapioca/run_gem_rbi_check"
+
+module Tapioca
+  module RubyLsp
+    class RunGemRbiCheckSpec < SpecWithProject
+      FOO_RB = "module Foo; end"
+
+      before(:all) do
+        @project = mock_project
+      end
+
+      describe "without git" do
+        before do
+          @project.bundle_install!
+        end
+
+        it "does nothing if there is no git repo" do
+          foo = mock_gem("foo", "0.0.1") do
+            write!("lib/foo.rb", FOO_RB)
+          end
+          @project.require_mock_gem(foo)
+
+          @project.bundle_install!
+          check = ::RubyLsp::Tapioca::RunGemRbiCheck.new
+          check.run(@project.absolute_path)
+
+          assert check.stdout.include?("Not a git repository")
+        end
+      end
+
+      describe "with git" do
+        before do
+          @project.write!("Gemfile", @project.tapioca_gemfile)
+          @project.bundle_install!
+          @project.exec("git init")
+          @project.exec("git add .")
+          @project.exec("git commit -m 'Initial commit'")
+        end
+
+        after do
+          @project.remove!("sorbet/rbi")
+          @project.remove!(".git")
+          @project.remove!("Gemfile")
+          @project.remove!("Gemfile.lock")
+        end
+
+        it "creates the RBI for a newly added gem" do
+          foo = mock_gem("foo", "0.0.1") do
+            write!("lib/foo.rb", FOO_RB)
+          end
+          @project.require_mock_gem(foo)
+          @project.bundle_install!
+
+          check = ::RubyLsp::Tapioca::RunGemRbiCheck.new
+          check.run(@project.absolute_path)
+
+          assert_project_file_exist("sorbet/rbi/gems/[email protected]")
+        end
+
+        it "regenerates RBI when a gem version changes" do
+          foo = mock_gem("foo", "0.0.1") do
+            write!("lib/foo.rb", FOO_RB)
+          end
+          @project.require_mock_gem(foo)
+          @project.bundle_install!
+
+          check = ::RubyLsp::Tapioca::RunGemRbiCheck.new
+          check.run(@project.absolute_path)
+
+          assert_project_file_exist("sorbet/rbi/gems/[email protected]")
+
+          # Modify the gem
+          foo.update("0.0.2")
+          @project.bundle_install!
+
+          check.run(@project.absolute_path)
+
+          assert_project_file_exist("sorbet/rbi/gems/[email protected]")
+        end
+
+        it "removes RBI file when a gem is removed" do
+          foo = mock_gem("foo", "0.0.1") do
+            write!("lib/foo.rb", FOO_RB)
+          end
+          @project.require_mock_gem(foo)
+          @project.bundle_install!
+
+          check1 = ::RubyLsp::Tapioca::RunGemRbiCheck.new
+          check1.run(@project.absolute_path)
+
+          assert_project_file_exist("sorbet/rbi/gems/[email protected]")
+
+          @project.exec("git restore Gemfile Gemfile.lock")
+
+          check2 = ::RubyLsp::Tapioca::RunGemRbiCheck.new
+          check2.run(@project.absolute_path)
+
+          refute_project_file_exist("sorbet/rbi/gems/[email protected]")
+        end
+
+        it "deletes untracked RBI files" do
+          @project.bundle_install!
+          FileUtils.mkdir_p("#{@project.absolute_path}/sorbet/rbi/gems")
+          # Create an untracked RBI file
+          FileUtils.touch("#{@project.absolute_path}/sorbet/rbi/gems/[email protected]")
+
+          assert_project_file_exist("/sorbet/rbi/gems/[email protected]")
+
+          check = ::RubyLsp::Tapioca::RunGemRbiCheck.new
+          check.run(@project.absolute_path)
+
+          refute_project_file_exist("sorbet/rbi/gems/[email protected]")
+        end
+
+        it "restores deleted RBI files" do
+          @project.bundle_install!
+          FileUtils.mkdir_p("#{@project.absolute_path}/sorbet/rbi/gems")
+          # Create and delete a tracked RBI file
+          FileUtils.touch("#{@project.absolute_path}/sorbet/rbi/gems/[email protected]")
+          @project.exec("git add sorbet/rbi/gems/[email protected]")
+          @project.exec("git commit -m 'Add foo RBI'")
+          FileUtils.rm("#{@project.absolute_path}/sorbet/rbi/gems/[email protected]")
+
+          refute_project_file_exist("sorbet/rbi/gems/[email protected]")
+
+          check = ::RubyLsp::Tapioca::RunGemRbiCheck.new
+          check.run(@project.absolute_path)
+
+          assert_project_file_exist("sorbet/rbi/gems/[email protected]")
+
+          # Clean-up commit
+          @project.exec("git reset --hard HEAD^")
+        end
+      end
+    end
+  end
+end