Support module declarations with self types

lib/rbs/inline/ast/declarations.rb

@@ -76,16 +76,28 @@ class ModuleDecl < Base
           attr_reader :node
           attr_reader :members
           attr_reader :comments
+          attr_reader :inner_comments
 
           def initialize(node, comments)
             @node = node
             @comments = comments
             @members = []
+            @inner_comments = []
           end
 
           def module_name
             type_name(node.constant_path)
           end
+
+          def module_selfs
+            inner_comments.flat_map do |comment|
+              comment.annotations.filter_map do |ann|
+                if ann.is_a?(AST::Annotations::ModuleSelf)
+                  ann
+                end
+              end
+            end
+          end
         end
       end
     end

lib/rbs/inline/parser.rb

@@ -92,6 +92,17 @@ def visit_module_node(node)
         push_class_module_decl(module_decl) do
           visit node.body
         end
+
+        annotation_range = node.location.start_line .. node.location.end_line
+        annotations = comments.each_value.select do |annotation|
+          range = annotation.line_range
+          annotation_range.begin < range.begin && range.end < annotation_range.end
+        end
+
+        annotations.each do |annot|
+          comments.delete(annot.line_range.end)
+          module_decl.inner_comments.push(annot)
+        end
       end
 
       def visit_def_node(node)

lib/rbs/inline/writer.rb

@@ -103,11 +103,13 @@ def translate_module_decl(decl)
           end
         end
 
+        self_types = decl.module_selfs.map { _1.constraint }.compact
+
         RBS::AST::Declarations::Module.new(
           name: decl.module_name,
           type_params: [],
           members: members,
-          self_types: [],
+          self_types: self_types,
           annotations: [],
           location: nil,
           comment: comment

sig/rbs/inline/ast/declarations.rbs

@@ -38,9 +38,13 @@ module RBS
 
           attr_reader comments: AnnotationParser::ParsingResult?
 
+          attr_reader inner_comments: Array[AnnotationParser::ParsingResult]
+
           def initialize: (Prism::ModuleNode, AnnotationParser::ParsingResult?) -> void
 
           %a{pure} def module_name: () -> TypeName?
+
+          %a{pure} def module_selfs: () -> Array[Annotations::ModuleSelf]
         end
       end
     end

test/rbs/inline/writer_test.rb

@@ -6,6 +6,7 @@ class RBS::Inline::WriterTest < Minitest::Test
   include RBS::Inline
 
   def translate(src)
+    src = "# rbs_inline: enabled\n\n" + src
     uses, decls = Parser.parse(Prism.parse(src, filepath: "a.rb"))
     Writer.write(uses, decls)
   end
@@ -348,4 +349,21 @@ def length: ...
       end
     RBS
   end
+
+  def test_module_self
+    output = translate(<<~RUBY)
+      module Foo
+        # @rbs module-self BasicObject
+
+        def foo
+        end
+      end
+    RUBY
+
+    assert_equal <<~RBS, output
+      module Foo : BasicObject
+        def foo: () -> untyped
+      end
+    RBS
+  end
 end