Implement StringIO#pread

Both for being closer to real IOs and also because it's a convenient API in multithreaded scenarios.
ruby · Jun 7, 2023 · a8d73d6 · a8d73d6
1 parent 1587d36
commit a8d73d6
Show file tree

Hide file tree

Showing 3 changed files with 124 additions and 0 deletions.
diff --git a/ext/java/org/jruby/ext/stringio/StringIO.java b/ext/java/org/jruby/ext/stringio/StringIO.java
@@ -907,6 +907,68 @@ public IRubyObject read(ThreadContext context, IRubyObject[] args) {
         return string;
     }
 
+    @JRubyMethod(name = "pread", required = 2, optional = 1)
+    public IRubyObject pread(ThreadContext context, IRubyObject[] args) {
+        checkReadable();
+
+        final Ruby runtime = context.runtime;
+        IRubyObject str = context.nil;
+        int len;
+        int offset;
+
+        StringIOData ptr = this.ptr;
+        final RubyString string;
+
+        switch (args.length) {
+            case 3:
+                str = args[2];
+                if (!str.isNil()) {
+                    str = str.convertToString();
+                    ((RubyString) str).modify();
+                }
+            case 2:
+                len = RubyNumeric.fix2int(args[0]);
+                offset = RubyNumeric.fix2int(args[1]);
+                if (!args[0].isNil()) {
+                    len = RubyNumeric.fix2int(args[0]);
+
+                    if (len < 0) {
+                        throw runtime.newArgumentError("negative length " + len + " given");
+                    }
+
+                    if (offset < 0) {
+                        throw runtime.newErrnoEINVALError("pread: Invalid offset argument");
+                    }
+                }
+                break;
+            default:
+                throw runtime.newArgumentError(args.length, 0, 2);
+        }
+
+        synchronized (ptr) {
+            if (offset >= ptr.string.size()) {
+                throw context.runtime.newEOFError();
+            }
+
+            if (str.isNil()) {
+                return strioSubstr(runtime, offset, len, ASCIIEncoding.INSTANCE);
+            }
+
+            string = (RubyString) str;
+            int rest = ptr.string.size() - offset;
+            if (len > rest) len = rest;
+            string.resize(len);
+            ByteList strByteList = string.getByteList();
+            byte[] strBytes = strByteList.getUnsafeBytes();
+            ByteList dataByteList = ptr.string.getByteList();
+            byte[] dataBytes = dataByteList.getUnsafeBytes();
+            System.arraycopy(dataBytes, dataByteList.getBegin() + offset, strBytes, strByteList.getBegin(), len);
+            string.setEncoding(ASCIIEncoding.INSTANCE);
+        }
+
+        return string;
+    }
+
     @JRubyMethod(name = "readlines")
     public IRubyObject readlines(ThreadContext context) {
         return Getline.getlineCall(context, GETLINE_ARY, this, getEncoding());

diff --git a/ext/stringio/stringio.c b/ext/stringio/stringio.c
@@ -1583,6 +1583,48 @@ strio_read(int argc, VALUE *argv, VALUE self)
     return str;
 }
 
+/*
+ *  call-seq:
+ *    pread(maxlen, offset)             -> string
+ *    pread(maxlen, offset, out_string) -> string
+ *
+ *  See IO#pread.
+ */
+static VALUE
+strio_pread(int argc, VALUE *argv, VALUE self)
+{
+    VALUE rb_len, rb_offset, rb_buf;
+    rb_scan_args(argc, argv, "21", &rb_len, &rb_offset, &rb_buf);
+    long len = NUM2LONG(rb_len);
+    long offset = NUM2LONG(rb_offset);
+
+    if (len < 0) {
+        rb_raise(rb_eArgError, "negative string size (or size too big): %" PRIsVALUE, rb_len);
+    }
+
+    if (offset < 0) {
+        rb_syserr_fail_str(EINVAL, rb_sprintf("pread: Invalid offset argument: %" PRIsVALUE, rb_offset));
+    }
+
+    struct StringIO *ptr = readable(self);
+
+    if (offset >= RSTRING_LEN(ptr->string)) {
+        rb_eof_error();
+    }
+
+    if (NIL_P(rb_buf)) {
+        return strio_substr(ptr, offset, len, rb_ascii8bit_encoding());
+    }
+
+    long rest = RSTRING_LEN(ptr->string) - offset;
+    if (len > rest) len = rest;
+    rb_str_resize(rb_buf, len);
+    rb_enc_associate(rb_buf, rb_ascii8bit_encoding());
+    MEMCPY(RSTRING_PTR(rb_buf), RSTRING_PTR(ptr->string) + offset, char, len);
+    return rb_buf;
+}
+
+
 /*
  * call-seq:
  *   strio.sysread(integer[, outbuf])    -> string
@@ -1843,6 +1885,7 @@ Init_stringio(void)
     rb_define_method(StringIO, "gets", strio_gets, -1);
     rb_define_method(StringIO, "readlines", strio_readlines, -1);
     rb_define_method(StringIO, "read", strio_read, -1);
+    rb_define_method(StringIO, "pread", strio_pread, -1);
 
     rb_define_method(StringIO, "write", strio_write_m, -1);
     rb_define_method(StringIO, "putc", strio_putc, 1);

diff --git a/test/stringio/test_stringio.rb b/test/stringio/test_stringio.rb
@@ -729,6 +729,25 @@ def test_sysread
     assert_equal Encoding::ASCII_8BIT, f.sysread(3).encoding
   end
 
+  def test_pread
+    f = StringIO.new("pread")
+    f.read
+
+    assert_equal "pre".b, f.pread(3, 0)
+    assert_equal "read".b, f.pread(4, 1)
+    assert_equal Encoding::ASCII_8BIT, f.pread(4, 1).encoding
+
+    buf = "".b
+    f.pread(3, 0, buf)
+    assert_equal "pre".b, buf
+    f.pread(4, 1, buf)
+    assert_equal "read".b, buf
+
+    assert_raise(EOFError) { f.pread(1, 5) }
+    assert_raise(ArgumentError) { f.pread(-1, 0) }
+    assert_raise(Errno::EINVAL) { f.pread(3, -1) }
+  end
+
   def test_size
     f = StringIO.new("1234")
     assert_equal(4, f.size)