Improve runtime test

lib/rbs/definition.rb

@@ -50,11 +50,13 @@ def update(type: self.type, member: self.member, defined_in: self.defined_in, im
       attr_reader :super_method
       attr_reader :defs
       attr_reader :accessibility
+      attr_reader :extra_annotations
 
-      def initialize(super_method:, defs:, accessibility:)
+      def initialize(super_method:, defs:, accessibility:, annotations: [])
         @super_method = super_method
         @defs = defs
         @accessibility = accessibility
+        @extra_annotations = annotations
       end
 
       def defined_in
@@ -74,7 +76,7 @@ def comments
       end
 
       def annotations
-        @annotations ||= defs.flat_map(&:annotations)
+        @annotations ||= @extra_annotations + defs.flat_map(&:annotations)
       end
 
       # @deprecated

lib/rbs/definition_builder.rb

@@ -934,7 +934,8 @@ def build_one_singleton(type_name)
                     implemented_in: nil
                   )
                 end,
-                accessibility: :public
+                accessibility: :public,
+                annotations: [AST::Annotation.new(location: nil, string: "rbs:test:target")]
               )
             end
           end

lib/rbs/test/hook.rb

@@ -7,6 +7,7 @@ module Hook
       OPERATORS = {
         :== => "eqeq",
         :=== => "eqeqeq",
+        :!= => "noteq",
         :+ => "plus",
         :- => "minus",
         :* => "star",

lib/rbs/test/tester.rb

@@ -17,6 +17,22 @@ def builder
         @builder ||= DefinitionBuilder.new(env: env)
       end
 
+      def skip_method?(type_name, method)
+        if method.implemented_in == type_name
+          if method.annotations.any? {|a| a.string == "rbs:test:skip" }
+            :skip
+          else
+            false
+          end
+        else
+          if method.annotations.any? {|a| a.string == "rbs:test:target" }
+            false
+          else
+            :implemented_in
+          end
+        end
+      end
+
       def install!(klass, sample_size:)
         RBS.logger.info { "Installing runtime type checker in #{klass}..." }
 
@@ -27,7 +43,11 @@ def install!(klass, sample_size:)
           Observer.register(instance_key, MethodCallTester.new(klass, builder, definition, kind: :instance, sample_size: sample_size))
 
           definition.methods.each do |name, method|
-            if method.implemented_in == type_name
+            if reason = skip_method?(type_name, method)
+              unless reason == :implemented_in
+                RBS.logger.info { "Skipping ##{name} because of `#{reason}`..." }
+              end
+            else
               RBS.logger.info { "Setting up method hook in ##{name}..." }
               Hook.hook_instance_method klass, name, key: instance_key
             end
@@ -39,7 +59,11 @@ def install!(klass, sample_size:)
           Observer.register(singleton_key, MethodCallTester.new(klass.singleton_class, builder, definition, kind: :singleton, sample_size: sample_size))
 
           definition.methods.each do |name, method|
-            if method.implemented_in == type_name || name == :new
+            if reason = skip_method?(type_name, method)
+              unless reason == :implemented_in
+                RBS.logger.info { "Skipping .#{name} because of `#{reason}`..." }
+              end
+            else
               RBS.logger.info { "Setting up method hook in .#{name}..." }
               Hook.hook_singleton_method klass, name, key: singleton_key
             end

lib/rbs/test/type_check.rb

@@ -5,12 +5,15 @@ class TypeCheck
       attr_reader :builder
       attr_reader :sample_size
 
+      attr_reader :const_cache
+
       DEFAULT_SAMPLE_SIZE = 100
 
       def initialize(self_class:, builder:, sample_size:)
         @self_class = self_class
         @builder = builder
         @sample_size = sample_size
+        @const_cache = {}
       end
 
       def overloaded_call(method, method_name, call, errors:)
@@ -179,8 +182,25 @@ def zip_args(args, fun, &block)
         end
       end
 
-      def sample(array)
-        sample_size && (array.size > sample_size) ? array.sample(sample_size) : array
+      def each_sample(array, &block)
+        if block
+          if sample_size && array.size > sample_size
+            if sample_size > 0
+              size = array.size
+              sample_size.times do
+                yield array[rand(size)]
+              end
+            end
+          else
+            array.each(&block)
+          end
+        else
+          enum_for :each_sample, array
+        end
+      end
+
+      def get_class(type_name)
+        const_cache[type_name] ||= Object.const_get(type_name.to_s)
       end
 
       def value(val, type)
@@ -204,17 +224,14 @@ def value(val, type)
         when Types::Bases::Instance
           Test.call(val, IS_AP, self_class)
         when Types::ClassInstance
-          klass = Object.const_get(type.name.to_s)
+          klass = get_class(type.name)
           case
           when klass == ::Array
-            Test.call(val, IS_AP, klass) && sample(val).yield_self do |val|
-              val.all? {|v| value(v, type.args[0]) }
-            end
+            Test.call(val, IS_AP, klass) && each_sample(val).all? {|v| value(v, type.args[0]) }
           when klass == ::Hash
-            Test.call(val, IS_AP, klass) && sample(val.keys).yield_self do |keys|
-              values = val.values_at(*keys)
-              keys.all? {|key| value(key, type.args[0]) } && values.all? {|v| value(v, type.args[1]) }
-            end 
+            Test.call(val, IS_AP, klass) && each_sample(val.keys).all? do |key|
+              value(key, type.args[0]) && value(val[key], type.args[1])
+            end
           when klass == ::Range
             Test.call(val, IS_AP, klass) && value(val.begin, type.args[0]) && value(val.end, type.args[0])
           when klass == ::Enumerator
@@ -235,7 +252,7 @@ def value(val, type)
                 end
               end
 
-              sample(values).all? do |v|
+              each_sample(values).all? do |v|
                 if v.size == 1
                   # Only one block argument.
                   value(v[0], type.args[0]) || value(v, type.args[0])
@@ -253,7 +270,7 @@ def value(val, type)
             Test.call(val, IS_AP, klass)
           end
         when Types::ClassSingleton
-          klass = Object.const_get(type.name.to_s)
+          klass = get_class(type.name)
           val == klass
         when Types::Interface
           methods = Set.new(Test.call(val, METHODS))

test/rbs/test/type_check_test.rb

@@ -81,20 +81,10 @@ def test_type_check_hash_sampling
 
         typecheck = Test::TypeCheck.new(self_class: Integer, builder: builder, sample_size: 100)
 
-        # hash = Array.new(100) {|i| [i, i.to_s] }.to_h
-
         assert typecheck.value({}, parse_type("::Hash[::Integer, ::String]"))
         assert typecheck.value(Array.new(100) {|i| [i, i.to_s] }.to_h, parse_type("::Hash[::Integer, ::String]"))
 
         assert typecheck.value(Array.new(1000) {|i| [i, i.to_s] }.to_h, parse_type("::Hash[::Integer, ::String]"))
-        refute typecheck.value(
-          Array.new(99) {|i| [i, i.to_s] }.to_h.merge({ foo: 'bar', bar: 'baz', baz: 'foo' }),
-          parse_type("::Hash[::Integer, ::String]")
-        )
-        refute typecheck.value(
-          Array.new(99) {|i| [i, i.to_s] }.to_h.merge({ 1001 => :bar, 1002 => :baz, 1003 => :foo }),
-          parse_type("::Hash[::Integer, ::String]")
-        )
       end
     end
   end

test/test_helper.rb

@@ -160,7 +160,7 @@ def assert_write(decls, string)
   def assert_sampling_check(builder, sample_size, array)
     checker = RBS::Test::TypeCheck.new(self_class: Integer, builder: builder, sample_size: sample_size)
 
-    sample = checker.sample(array)
+    sample = checker.each_sample(array).to_a
 
     assert_operator(sample.size, :<=, array.size)
     assert_operator(sample.size, :<=, sample_size) unless sample_size.nil?