From 001e75be16354b8239985ee12c79b3df1f2aae9e Mon Sep 17 00:00:00 2001
From: Alexandre Terrasa <alexandre.terrasa@shopify.com>
Date: Wed, 18 Dec 2024 12:11:05 -0500
Subject: [PATCH] Add command to strip Sorbet signatures from files

Signed-off-by: Alexandre Terrasa <alexandre.terrasa@shopify.com>
---
 lib/spoom/cli/srb/sigs.rb       | 16 +++++++++--
 lib/spoom/sorbet/sigs.rb        | 12 ++++++++
 test/spoom/cli/srb/sigs_test.rb | 33 ++++++++++++++++++++++
 test/spoom/sorbet/sigs_test.rb  | 49 +++++++++++++++++++++++++++++++++
 4 files changed, 107 insertions(+), 3 deletions(-)

diff --git a/lib/spoom/cli/srb/sigs.rb b/lib/spoom/cli/srb/sigs.rb
index 938eae3e..bf1261b4 100644
--- a/lib/spoom/cli/srb/sigs.rb
+++ b/lib/spoom/cli/srb/sigs.rb
@@ -29,7 +29,16 @@ def translate(*paths)
         end
 
         desc "strip", "Strip all the signatures from the files"
-        def strip(*path)
+        def strip(*paths)
+          files = collect_files(paths)
+
+          say("Stripping signatures from `#{files.size}` file#{files.size == 1 ? "" : "s"}...\n\n")
+
+          transformed_files = transform_files(files) do |_file, contents|
+            Spoom::Sorbet::Sigs.strip(contents)
+          end
+
+          say("Stripped signatures from `#{transformed_files}` file#{transformed_files == 1 ? "" : "s"}.")
         end
 
         no_commands do
@@ -57,9 +66,10 @@ def transform_files(files, &block)
 
             files.each do |file|
               contents = File.read(file)
+              first_line = contents.lines.first
 
-              if contents.lines.first&.start_with?("# encoding:")
-                encoding = T.must(contents.lines.first).gsub(/^#\s*encoding:\s*/, "").strip
+              if first_line&.start_with?("# encoding:")
+                encoding = T.must(first_line).gsub(/^#\s*encoding:\s*/, "").strip
                 contents = contents.force_encoding(encoding)
               end
 
diff --git a/lib/spoom/sorbet/sigs.rb b/lib/spoom/sorbet/sigs.rb
index 1473df9f..2d055d58 100644
--- a/lib/spoom/sorbet/sigs.rb
+++ b/lib/spoom/sorbet/sigs.rb
@@ -9,6 +9,18 @@ class Sigs
       class << self
         extend T::Sig
 
+        sig { params(ruby_contents: String).returns(String) }
+        def strip(ruby_contents)
+          sigs = collect_sigs(ruby_contents)
+          lines_to_strip = sigs.flat_map { |sig, _| (sig.loc&.begin_line..sig.loc&.end_line).to_a }
+
+          lines = []
+          ruby_contents.lines.each_with_index do |line, index|
+            lines << line unless lines_to_strip.include?(index + 1)
+          end
+          lines.join
+        end
+
         sig { params(ruby_contents: String).returns(String) }
         def rbi_to_rbs(ruby_contents)
           ruby_contents = ruby_contents.dup
diff --git a/test/spoom/cli/srb/sigs_test.rb b/test/spoom/cli/srb/sigs_test.rb
index 25b3ee2e..50dfbacd 100644
--- a/test/spoom/cli/srb/sigs_test.rb
+++ b/test/spoom/cli/srb/sigs_test.rb
@@ -11,6 +11,39 @@ def setup
           @project.bundle_install!
         end
 
+        # strip
+
+        def test_strip_sigs
+          @project.write!("a/file1.rb", <<~RB)
+            sig { void }
+            def foo; end
+
+            class B
+              sig { void }
+              def bar; end
+            end
+          RB
+
+          result = @project.spoom("srb sigs strip --no-color")
+
+          assert_equal(<<~OUT, result.out)
+            Stripping signatures from `1` file...
+
+            Stripped signatures from `1` file.
+          OUT
+          assert(result.status)
+
+          assert_equal(<<~RB, @project.read("a/file1.rb"))
+            def foo; end
+
+            class B
+              def bar; end
+            end
+          RB
+        end
+
+        # translate
+
         def test_only_supports_translation_from_rbi
           result = @project.spoom("srb sigs translate --from rbs")
 
diff --git a/test/spoom/sorbet/sigs_test.rb b/test/spoom/sorbet/sigs_test.rb
index 4c096b5c..b9d4aa35 100644
--- a/test/spoom/sorbet/sigs_test.rb
+++ b/test/spoom/sorbet/sigs_test.rb
@@ -6,6 +6,55 @@
 module Spoom
   module Sorbet
     class SigsTest < Minitest::Test
+      # strip
+
+      def test_strip_empty
+        contents = ""
+        assert_equal(contents, Sigs.strip(contents))
+      end
+
+      def test_strip_no_sigs
+        contents = <<~RB
+          class A
+            def foo; end
+          end
+        RB
+
+        assert_equal(contents, Sigs.strip(contents))
+      end
+
+      def test_strip_sigs
+        contents = <<~RB
+          class A
+            sig { returns(Integer) }
+            attr_accessor :a
+
+            sig { void }
+            def foo; end
+
+            module B
+              sig { void }
+              sig { returns(Integer) }
+              def bar; end
+            end
+          end
+        RB
+
+        assert_equal(<<~RB, Sigs.strip(contents))
+          class A
+            attr_accessor :a
+
+            def foo; end
+
+            module B
+              def bar; end
+            end
+          end
+        RB
+      end
+
+      # translate
+
       def test_translate_empty
         contents = ""
         assert_equal(contents, Sigs.rbi_to_rbs(contents))