Module#define_method supports proc argument

2 files changed, 25 insertions, 1 deletions

diff --git a/src/class.c b/src/class.c
index 47a6c846b..00fc9c12e 100644
--- a/src/class.c
+++ b/src/class.c
@@ -1773,9 +1773,21 @@ mod_define_method(mrb_state *mrb, mrb_value self)
   struct RClass *c = mrb_class_ptr(self);
   struct RProc *p;
   mrb_sym mid;
+  mrb_value proc = mrb_undef_value();
   mrb_value blk;
 
-  mrb_get_args(mrb, "n&", &mid, &blk);
+  mrb_get_args(mrb, "n|o&", &mid, &proc, &blk);
+  switch (mrb_type(proc)) {
+    case MRB_TT_PROC:
+      blk = proc;
+      break;
+    case MRB_TT_UNDEF:
+      /* ignored */
+      break;
+    default:
+      mrb_raisef(mrb, E_TYPE_ERROR, "wrong argument type %S (expected Proc)", mrb_obj_value(mrb_obj_class(mrb, proc)));
+      break;
+  }
   if (mrb_nil_p(blk)) {
     mrb_raise(mrb, E_ARGUMENT_ERROR, "no block given");
   }
diff --git a/test/t/module.rb b/test/t/module.rb
index 4bde20fbe..6b0632414 100644
--- a/test/t/module.rb
+++ b/test/t/module.rb
@@ -494,6 +494,18 @@ end
 
 # Not ISO specified
 
+assert('Module#define_method') do
+  c = Class.new {
+    define_method(:m1) { :ok }
+    define_method(:m2, Proc.new { :ok })
+  }
+  assert_equal c.new.m1, :ok
+  assert_equal c.new.m2, :ok
+  assert_raise(TypeError) do
+    Class.new { define_method(:n1, nil) }
+  end
+end
+
 # @!group prepend
   assert('Module#prepend') do
     module M0