Merge pull request #350 from Shopify/at-rewriter-flatten-visibility

Introduce FlattenVisibilities rewriter

lib/rbi.rb

@@ -20,6 +20,7 @@ class Error < StandardError
 require "rbi/rewriters/deannotate"
 require "rbi/rewriters/filter_versions"
 require "rbi/rewriters/flatten_singleton_methods"
+require "rbi/rewriters/flatten_visibilities"
 require "rbi/rewriters/merge_trees"
 require "rbi/rewriters/nest_singleton_methods"
 require "rbi/rewriters/nest_non_public_methods"

lib/rbi/rewriters/flatten_visibilities.rb

@@ -0,0 +1,65 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RBI
+  module Rewriters
+    # Flattens visibility nodes into method nodes
+    #
+    # Example:
+    # ~~~rb
+    # class A
+    #   def m1; end
+    #   private
+    #   def m2; end
+    #   def m3; end
+    # end
+    # ~~~
+    #
+    # will be transformed into:
+    #
+    # ~~~rb
+    # class A
+    #   def m1; end
+    #   private def m2; end
+    #   private def m3; end
+    # end
+    # ~~~
+    class FlattenVisibilities < Visitor
+      extend T::Sig
+
+      sig { void }
+      def initialize
+        super
+
+        @current_visibility = T.let([Public.new], T::Array[Visibility])
+      end
+
+      sig { override.params(node: T.nilable(Node)).void }
+      def visit(node)
+        return unless node
+
+        case node
+        when Public, Protected, Private
+          @current_visibility[-1] = node
+          node.detach
+        when Tree
+          @current_visibility << Public.new
+          visit_all(node.nodes.dup)
+          @current_visibility.pop
+        when Attr, Method
+          node.visibility = T.must(@current_visibility.last)
+        end
+      end
+    end
+  end
+
+  class Tree
+    extend T::Sig
+
+    sig { void }
+    def flatten_visibilities!
+      visitor = Rewriters::FlattenVisibilities.new
+      visitor.visit(self)
+    end
+  end
+end

test/rbi/rewriters/flatten_visibilities_test.rb

@@ -0,0 +1,107 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+module RBI
+  class FlattenVisibilitiesTest < Minitest::Test
+    include TestHelper
+
+    def test_flatten_visibilities_in_tree
+      tree = parse_rbi(<<~RBI)
+        def m1; end
+        def m2; end
+        protected
+        def m3; end
+        def m4; end
+        private
+        def m5; end
+        def m6; end
+      RBI
+
+      tree.flatten_visibilities!
+
+      assert_equal(<<~RBI, tree.string)
+        def m1; end
+        def m2; end
+        protected def m3; end
+        protected def m4; end
+        private def m5; end
+        private def m6; end
+      RBI
+    end
+
+    def test_flatten_visibilityies_in_scopes
+      tree = parse_rbi(<<~RBI)
+        module S1
+          class S2
+            class << self
+              def m1; end
+
+              protected
+
+              def m2; end
+              attr_reader :m3
+            end
+
+            def m4; end
+
+            private
+
+            def m5; end
+            attr_writer m6
+          end
+
+          def m7; end
+
+          protected
+
+          def m8; end
+        end
+      RBI
+
+      tree.flatten_visibilities!
+
+      assert_equal(<<~RBI, tree.string)
+        module S1
+          class S2
+            class << self
+              def m1; end
+              protected def m2; end
+              protected attr_reader :m3
+            end
+
+            def m4; end
+            private def m5; end
+            private attr_writer :m6
+          end
+
+          def m7; end
+          protected def m8; end
+        end
+      RBI
+    end
+
+    def test_flatten_visibilities_for_singleton_methods
+      tree = parse_rbi(<<~RBI)
+        def self.m1; end
+
+        protected
+
+        def self.m2; end
+
+        private
+
+        def self.m3; end
+      RBI
+
+      tree.flatten_visibilities!
+
+      assert_equal(<<~RBI, tree.string)
+        def self.m1; end
+        protected def self.m2; end
+        private def self.m3; end
+      RBI
+    end
+  end
+end