lib/safety_check.rb in safety_check-0.0.1 vs lib/safety_check.rb in safety_check-0.0.2

@@ -13,27 +13,29 @@
         raise NoMethodError, "undefined method '#{method_name}' for #{receiver.inspect}"
       end
 
       unsafe_method_name = "#{method_name.to_s}_without_type_safety".to_sym
 
-      if arg_types.length != receiver.instance_method(method_name).arity
-        raise ArgumentError, "wrong number of arguments for #{method_name.inspect} (#{arg_types.length} for #{receiver.instance_method(method_name).arity})"
+      required_arity = receiver.instance_method(method_name).parameters.select { |type, name| type == :req }.length
+
+      if arg_types.length < required_arity
+        raise ArgumentError, "wrong number of arguments for #{method_name.inspect} (#{arg_types.length} for #{required_arity})"
       end
 
       receiver.send(:alias_method, unsafe_method_name, method_name)
 
-      receiver.send(:define_method, method_name) do |*args|
-        if args.length != arg_types.length
+      receiver.send(:define_method, method_name) do |*args, &block|
+        if args.length < required_arity
           raise ArgumentError, "wrong number of arguments (#{args.length} for #{arg_types.length}) when calling #{method_name.inspect}"
         end
 
-        arg_types.zip(args).each do |arg_type, arg|
+        args.zip(arg_types).each do |arg, arg_type|
           unless arg.is_a? arg_type
             raise ArgumentError, "expected #{arg.inspect} to be a #{arg_type} when calling #{method_name.inspect}"
           end
         end
 
-        self.send(unsafe_method_name, *args)
+        self.send(unsafe_method_name, *args, &block)
       end
     end
   end
 end