allow IO(-ish) objects as descriptor arguments

This should be more convenient to use than having
to remember to call "fileno" every time...

2 files changed, 98 insertions, 5 deletions

diff --git a/ext/io_splice/io_splice_ext.c b/ext/io_splice/io_splice_ext.c
index 3a1cf7a..6b27b59 100644
--- a/ext/io_splice/io_splice_ext.c
+++ b/ext/io_splice/io_splice_ext.c
@@ -10,6 +10,37 @@
 #include <limits.h>
 #include <alloca.h>
 
+#if ! HAVE_RB_IO_T
+#  define rb_io_t OpenFile
+#endif
+
+#ifdef GetReadFile
+#  define FPTR_TO_FD(fptr) (fileno(GetReadFile(fptr)))
+#else
+#  if !HAVE_RB_IO_T || (RUBY_VERSION_MAJOR == 1 && RUBY_VERSION_MINOR == 8)
+#    define FPTR_TO_FD(fptr) fileno(fptr->f)
+#  else
+#    define FPTR_TO_FD(fptr) fptr->fd
+#  endif
+#endif
+
+static int my_fileno(VALUE io)
+{
+        rb_io_t *fptr;
+
+        for (;;) {
+                switch (TYPE(io)) {
+                case T_FIXNUM: return NUM2INT(io);
+                case T_FILE: {
+                        GetOpenFile(io, fptr);
+                        return FPTR_TO_FD(fptr);
+                }
+                default:
+                        io = rb_convert_type(io, T_FILE, "IO", "to_io");
+                        /* retry */
+                }
+        }
+}
 #ifndef HAVE_RB_THREAD_BLOCKING_REGION
 /* partial emulation of the 1.9 rb_thread_blocking_region under 1.8 */
 #  include <rubysig.h>
@@ -123,8 +154,8 @@ static VALUE my_splice(VALUE self,
         struct splice_args a = {
                 .off_in = NIL_P(off_in) ? NULL : (i = NUM2OFFT(off_in), &i),
                 .off_out = NIL_P(off_out) ? NULL : (o = NUM2OFFT(off_out), &o),
-                .fd_in = NUM2INT(fd_in),
-                .fd_out = NUM2INT(fd_out),
+                .fd_in = my_fileno(fd_in),
+                .fd_out = my_fileno(fd_out),
                 .len = (size_t)NUM2ULONG(len),
                 .flags = NUM2UINT(flags),
         };
@@ -177,8 +208,8 @@ static VALUE my_tee(VALUE self,
 {
         long n;
         struct tee_args a = {
-                .fd_in = NUM2INT(fd_in),
-                .fd_out = NUM2INT(fd_out),
+                .fd_in = my_fileno(fd_in),
+                .fd_out = my_fileno(fd_out),
                 .len = (size_t)NUM2ULONG(len),
                 .flags = NUM2UINT(flags),
         };
@@ -278,7 +309,7 @@ static VALUE my_vmsplice(VALUE self, VALUE fd, VALUE data, VALUE flags)
         struct vmsplice_args a;
 
         ARY2IOVEC(a.iov, a.nr_segs, left, data);
-        a.fd = NUM2INT(fd);
+        a.fd = my_fileno(fd);
         a.flags = NUM2UINT(flags);
 
         for (;;) {
diff --git a/test/test_io_splice.rb b/test/test_io_splice.rb
index 142586a..e975c06 100644
--- a/test/test_io_splice.rb
+++ b/test/test_io_splice.rb
@@ -22,6 +22,42 @@ class Test_IO_Splice < Test::Unit::TestCase
     assert_equal str, rd.sysread(size)
   end
 
+  def test_splice_io
+    str = 'abcde'
+    size = 5
+    rd, wr = IO.pipe
+    tmp = Tempfile.new('ruby_io_splice')
+
+    assert_nothing_raised {
+      tmp.syswrite(str)
+      tmp.sysseek(0)
+    }
+
+    nr = IO.splice(tmp, nil, wr, nil, size, 0)
+    assert_equal size, nr
+    assert_equal str, rd.sysread(size)
+  end
+
+  def test_splice_io_ish
+    str = 'abcde'
+    size = 5
+    rd, wr = IO.pipe
+    tmp = Tempfile.new('ruby_io_splice')
+    io_ish = [ tmp ]
+    def io_ish.to_io
+      first.to_io
+    end
+
+    assert_nothing_raised {
+      tmp.syswrite(str)
+      tmp.sysseek(0)
+    }
+
+    nr = IO.splice(io_ish, nil, wr, nil, size, 0)
+    assert_equal size, nr
+    assert_equal str, rd.sysread(size)
+  end
+
   def test_splice_in_offset
     str = 'abcde'
     off = 3
@@ -101,6 +137,19 @@ class Test_IO_Splice < Test::Unit::TestCase
     }
   end
 
+  def test_tee_io
+    str = 'abcde'
+    size = 5
+    rda, wra = IO.pipe
+    rdb, wrb = IO.pipe
+
+    assert_nothing_raised { wra.syswrite(str) }
+    nr = IO.tee(rda, wrb, size, 0)
+    assert_equal 5, nr
+    assert_equal str, rdb.sysread(5)
+    assert_equal str, rda.sysread(5)
+  end
+
   def test_vmsplice_array
     data = %w(hello world how are you today)
     r, w = IO.pipe
@@ -109,6 +158,14 @@ class Test_IO_Splice < Test::Unit::TestCase
     assert_equal data.join(''), r.readpartial(16384)
   end
 
+  def test_vmsplice_array_io
+    data = %w(hello world how are you today)
+    r, w = IO.pipe
+    n = IO.vmsplice(w, data, 0)
+    assert_equal data.join('').size, n
+    assert_equal data.join(''), r.readpartial(16384)
+  end
+
   def test_vmsplice_nonblock
     data = %w(hello world how are you today)
     r, w = IO.pipe
@@ -147,6 +204,11 @@ class Test_IO_Splice < Test::Unit::TestCase
     end
   end
 
+  def test_vmsplice_nil
+    data = %w(hello world how are you today)
+    assert_raises(TypeError) { IO.vmsplice(nil, data, 0) }
+  end
+
   def test_constants
     assert IO::Splice::PIPE_BUF > 0
     %w(move nonblock more gift).each { |x|