Add resolve-type-names magic comment

lib/rbs.rb

@@ -9,6 +9,7 @@
 require "ripper"
 require "logger"
 require "tsort"
+require "strscan"
 
 require "rbs/errors"
 require "rbs/buffer"

lib/rbs/ast/directives.rb

@@ -34,6 +34,16 @@ def initialize(clauses:, location:)
         end
       end
 
+      class ResolveTypeNames < Base
+        attr_reader :location
+
+        attr_reader :value
+
+        def initialize(value:, location:)
+          @value = value
+          @location = location
+        end
+      end
     end
   end
 end

lib/rbs/environment.rb

@@ -486,6 +486,28 @@ def validate_type_params
       end
     end
 
+    def resolve_signature(resolver, table, dirs, decls, only: nil)
+      map = UseMap.new(table: table)
+      dirs.each do |dir|
+        case dir
+        when AST::Directives::Use
+          dir.clauses.each do |clause|
+            map.build_map(clause)
+          end
+        end
+      end
+
+      decls = decls.map do |decl|
+        if only && !only.member?(decl)
+          decl
+        else
+          resolve_declaration(resolver, map, decl, outer: [], prefix: Namespace.root)
+        end
+      end
+
+      [dirs, decls]
+    end
+
     def resolve_type_names(only: nil)
       resolver = Resolver::TypeNameResolver.new(self)
       env = Environment.new
@@ -498,21 +520,10 @@ def resolve_type_names(only: nil)
       table.compute_children
 
       signatures.each do |buffer, (dirs, decls)|
-        map = UseMap.new(table: table)
-        dirs.each do |dir|
-          dir.clauses.each do |clause|
-            map.build_map(clause)
-          end
+        resolve = dirs.find { _1.is_a?(AST::Directives::ResolveTypeNames) } #: AST::Directives::ResolveTypeNames?
+        if !resolve || resolve.value
+          _, decls = resolve_signature(resolver, table, dirs, decls)
         end
-
-        decls = decls.map do |decl|
-          if only && !only.member?(decl)
-            decl
-          else
-            resolve_declaration(resolver, map, decl, outer: [], prefix: Namespace.root)
-          end
-        end
-
         env.add_signature(buffer: buffer, directives: dirs, decls: decls)
       end
 

lib/rbs/locator.rb

@@ -40,6 +40,8 @@ def find2(line:, column:)
     end
 
     def find_in_directive(pos, dir, array)
+      return false unless dir.is_a?(AST::Directives::Use)
+
       if test_loc(pos, location: dir.location)
         array.unshift(dir)
 

lib/rbs/parser_aux.rb

@@ -17,11 +17,48 @@ def self.parse_method_type(source, range: 0..., variables: [], require_eof: fals
 
     def self.parse_signature(source)
       buf = buffer(source)
-      dirs, decls = _parse_signature(buf, 0, buf.last_position)
+
+      resolved = magic_comment(buf)
+      start_pos =
+        if resolved
+          (resolved.location || raise).end_pos
+        else
+          0
+        end
+      dirs, decls = _parse_signature(buf, start_pos, buf.last_position)
+
+      if resolved
+        dirs = dirs.dup if dirs.frozen?
+        dirs.unshift(resolved)
+      end
 
       [buf, dirs, decls]
     end
 
+    def self.magic_comment(buf)
+      start_pos = 0
+
+      while true
+        case
+        when match = /\A#\s*(?<keyword>resolve-type-names)\s*(?<colon>:)\s+(?<value>true|false)$/.match(buf.content, start_pos)
+          value = match[:value] or raise
+
+          kw_offset = match.offset(:keyword) #: [Integer, Integer]
+          colon_offset = match.offset(:colon) #: [Integer, Integer]
+          value_offset = match.offset(:value) #: [Integer, Integer]
+
+          location = Location.new(buf, kw_offset[0], value_offset[1])
+          location.add_required_child(:keyword, kw_offset[0]...kw_offset[1])
+          location.add_required_child(:colon, colon_offset[0]...colon_offset[1])
+          location.add_required_child(:value, value_offset[0]...value_offset[1])
+
+          return AST::Directives::ResolveTypeNames.new(value: value == "true", location: location)
+        else
+          return
+        end
+      end
+    end
+
     def self.lex(source)
       buf = buffer(source)
       list = _lex(buf, buf.last_position)

lib/rbs/writer.rb

@@ -77,14 +77,19 @@ def write_comment(comment)
     end
 
     def write(contents)
-      dirs = contents.select {|c| c.is_a?(AST::Directives::Base) } #: Array[AST::Directives::t]
+      resolves = contents.select { _1.is_a?(AST::Directives::ResolveTypeNames) } #: Array[AST::Directives::ResolveTypeNames]
+      uses = contents.select {|c| c.is_a?(AST::Directives::Use) } #: Array[AST::Directives::Use]
       decls = contents.select {|c| c.is_a?(AST::Declarations::Base) } #: Array[AST::Declarations::t]
 
-      dirs.each do |dir|
-        write_directive(dir)
+      if first_resolves = resolves.first
+        puts "# resolve-type-names: #{first_resolves.value}"
+        puts
       end
 
-      puts unless dirs.empty?
+      uses.each do |dir|
+        write_use_directive(dir)
+      end
+      puts unless uses.empty?
 
       [nil, *decls].each_cons(2) do |prev, decl|
         raise unless decl
@@ -94,7 +99,7 @@ def write(contents)
       end
     end
 
-    def write_directive(dir)
+    def write_use_directive(dir)
       clauses = dir.clauses.map do |clause|
         case clause
         when AST::Directives::Use::SingleClause

sig/directives.rbs

@@ -1,7 +1,7 @@
 module RBS
   module AST
     module Directives
-      type t = Use
+      type t = Use | ResolveTypeNames
 
       class Base
       end
@@ -56,6 +56,22 @@ module RBS
 
         def initialize: (clauses: Array[clause], location: loc?) -> void
       end
+
+      class ResolveTypeNames < Base
+        # ```
+        # # resolve-type-names: false
+        #   ^^^^^^^^^^^^^^^^^^          keyword
+        #                     ^         colon
+        #                       ^^^^^   value
+        # ```
+        type loc = Location[:keyword | :colon | :value, bot]
+
+        attr_reader location: loc?
+
+        attr_reader value: bool
+
+        def initialize: (value: bool, location: loc?) -> void
+      end
     end
   end
 end

sig/environment.rbs

@@ -141,6 +141,8 @@ module RBS
     #
     def resolve_type_names: (?only: Set[AST::Declarations::t]?) -> Environment
 
+    def resolve_signature: (Resolver::TypeNameResolver, UseMap::Table, Array[AST::Directives::t], Array[AST::Declarations::t], ?only: Set[AST::Declarations::t]?) -> [Array[AST::Directives::t], Array[AST::Declarations::t]]
+
     def inspect: () -> String
 
     def buffers: () -> Array[Buffer]

sig/parser.rbs

@@ -68,6 +68,10 @@ module RBS
     #
     def self.parse_signature: (Buffer | String) -> [Buffer, Array[AST::Directives::t], Array[AST::Declarations::t]]
 
+    # Returns the magic comment from the buffer
+    #
+    def self.magic_comment: (Buffer) -> AST::Directives::ResolveTypeNames?
+
     # Parse whole RBS file and return result.
     #
     # ```ruby
@@ -86,7 +90,7 @@ module RBS
 
     def self._parse_method_type: (Buffer, Integer start_pos, Integer end_pos, Array[Symbol] variables, bool require_eof) -> MethodType?
 
-    def self._parse_signature: (Buffer, Integer end_pos) -> [Array[AST::Directives::t], Array[AST::Declarations::t]]
+    def self._parse_signature: (Buffer, Integer start_pos, Integer end_pos) -> [Array[AST::Directives::t], Array[AST::Declarations::t]]
 
     def self._lex: (Buffer, Integer end_pos) -> Array[[Symbol, Location[untyped, untyped]]]
 

sig/writer.rbs

@@ -96,7 +96,7 @@ module RBS
 
     def format_annotation: (AST::Annotation) -> String
 
-    def write_directive: (AST::Directives::t) -> void
+    def write_use_directive: (AST::Directives::Use) -> void
 
     def write_annotation: (Array[AST::Annotation]) -> void
 

test/rbs/environment_test.rb

@@ -549,4 +549,40 @@ class OB
       assert_operator env.class_decls, :key?, RBS::TypeName.parse("::OB")
     end
   end
+
+  def test_resolve_type_names_magic_comment
+    buf, dirs, decls = RBS::Parser.parse_signature(<<-RBS)
+# resolve-type-names: false
+
+type t = s
+
+type s = untyped
+    RBS
+
+    env = Environment.new
+    env.add_signature(buffer: buf, directives: dirs, decls: decls)
+
+    env.resolve_type_names.tap do |env|
+      alias_decl = env.type_alias_decls[RBS::TypeName.parse("::t")]
+      assert_equal "s", alias_decl.decl.type.to_s
+    end
+  end
+
+  def test_resolve_type_names_magic_comment__true
+    buf, dirs, decls = RBS::Parser.parse_signature(<<-RBS)
+# resolve-type-names: true
+
+type t = s
+
+type s = untyped
+    RBS
+
+    env = Environment.new
+    env.add_signature(buffer: buf, directives: dirs, decls: decls)
+
+    env.resolve_type_names.tap do |env|
+      alias_decl = env.type_alias_decls[RBS::TypeName.parse("::t")]
+      assert_equal "::s", alias_decl.decl.type.to_s
+    end
+  end
 end

test/rbs/signature_parsing_test.rb

@@ -2171,4 +2171,53 @@ module Baz
       RBS
     end
   end
+
+  def test_resolved_directive
+    Parser.parse_signature(<<~RBS).tap do |_, dirs, _|
+        # resolve-type-names: false
+
+        module Baz
+        end
+      RBS
+
+      assert_equal 1, dirs.size
+      dirs[0].tap do
+        assert_instance_of RBS::AST::Directives::ResolveTypeNames, _1
+        assert_false _1.value
+        assert_equal "resolve-type-names: false", _1.location.source
+        assert_equal "resolve-type-names", _1.location[:keyword].source
+        assert_equal ":", _1.location[:colon].source
+        assert_equal "false", _1.location[:value].source
+      end
+    end
+
+    Parser.parse_signature(<<~RBS).tap do |_, dirs, _|
+        #resolve-type-names : true
+
+        module Baz
+        end
+      RBS
+
+      assert_equal 1, dirs.size
+      dirs[0].tap do
+        assert_instance_of RBS::AST::Directives::ResolveTypeNames, _1
+        assert_true _1.value
+        assert_equal "resolve-type-names : true", _1.location.source
+        assert_equal "resolve-type-names", _1.location[:keyword].source
+        assert_equal ":", _1.location[:colon].source
+        assert_equal "true", _1.location[:value].source
+      end
+    end
+
+    Parser.parse_signature(<<~RBS).tap do |_, dirs, _|
+        # This is a comment
+        # resolve-type-names: false
+
+        module Baz
+        end
+      RBS
+
+      assert_empty dirs
+    end
+  end
 end

test/rbs/writer_test.rb

@@ -328,6 +328,55 @@ def test___todo__
     assert_writer <<-SIG
 class Foo
   attr_reader name: __todo__
+end
+    SIG
+  end
+
+  def test_magic_comment
+    assert_writer <<-SIG
+# resolve-type-names: false
+
+class Foo
+end
+    SIG
+
+    assert_writer <<-SIG
+# resolve-type-names: false
+
+use String as S
+
+type s = S
+    SIG
+  end
+
+  def test_magic_comment_with_additional_new_line
+    result = format(<<-SIG)
+# resolve-type-names: false
+# Class Foo is something...
+class Foo
+end
+    SIG
+
+    assert_equal <<-SIG, result
+# resolve-type-names: false
+
+# Class Foo is something...
+class Foo
+end
+    SIG
+  end
+
+  def test_magic_comment_with_additional_new_line2
+    result = format(<<-SIG)
+# resolve-type-names: false
+class Foo
+end
+    SIG
+
+    assert_equal <<-SIG, result
+# resolve-type-names: false
+
+class Foo
 end
     SIG
   end