diff --git a/lib/openai/internal/util.rb b/lib/openai/internal/util.rb index e254972b..dff2080d 100644 --- a/lib/openai/internal/util.rb +++ b/lib/openai/internal/util.rb @@ -421,18 +421,38 @@ def close # # @param max_len [Integer, nil] # - # @return [String] + # @return [String, nil] private def read_enum(max_len) case max_len in nil - @stream.to_a.join + # `loop` rescues StopIteration, but this method handles it below. + # rubocop:disable Style/InfiniteLoop + @buf << @stream.next.b while true + # rubocop:enable Style/InfiniteLoop + in Integer if max_len.negative? + raise ArgumentError, "negative length #{max_len} given" in Integer - @buf << @stream.next while @buf.length < max_len - @buf.slice!(..max_len) + @buf << @stream.next.b while @buf.bytesize < max_len + read_buffer(max_len) end rescue StopIteration + return @buf.slice!(0..) if max_len.nil? + @stream = nil - @buf.slice!(0..) + return nil if @buf.bytesize.zero? + + read_buffer(max_len) + end + + # @api private + # + # @param max_len [Integer] + # + # @return [String] + private def read_buffer(max_len) + read = @buf.byteslice(0, max_len) + @buf = @buf.byteslice(max_len..) || String.new + read end # @api private @@ -444,14 +464,24 @@ def close def read(max_len = nil, out_string = nil) case @stream in nil - nil + raise ArgumentError, "negative length #{max_len} given" if max_len&.negative? + + read = max_len.nil? || max_len.zero? ? +"" : nil + case out_string + in String + out_string.replace(read || +"") + read.nil? ? nil : out_string + in nil + read + end in IO | StringIO @stream.read(max_len, out_string) in Enumerator read = read_enum(max_len) case out_string in String - out_string.replace(read) + out_string.replace(read || +"") + read.nil? ? nil : out_string in nil read end diff --git a/rbi/openai/internal/util.rbi b/rbi/openai/internal/util.rbi index f80846be..5b26f470 100644 --- a/rbi/openai/internal/util.rbi +++ b/rbi/openai/internal/util.rbi @@ -268,10 +268,15 @@ module OpenAI end # @api private - sig { params(max_len: T.nilable(Integer)).returns(String) } + sig { params(max_len: T.nilable(Integer)).returns(T.nilable(String)) } private def read_enum(max_len) end + # @api private + sig { params(max_len: Integer).returns(String) } + private def read_buffer(max_len) + end + # @api private sig do params( diff --git a/sig/openai/internal/util.rbs b/sig/openai/internal/util.rbs index 49cc4467..1c9f3f29 100644 --- a/sig/openai/internal/util.rbs +++ b/sig/openai/internal/util.rbs @@ -92,7 +92,9 @@ module OpenAI def close: -> void - private def read_enum: (Integer? max_len) -> String + private def read_enum: (Integer? max_len) -> String? + + private def read_buffer: (Integer max_len) -> String def read: (?Integer? max_len, ?String? out_string) -> String? diff --git a/test/openai/internal/util_test.rb b/test/openai/internal/util_test.rb index 560d786d..c989aa1f 100644 --- a/test/openai/internal/util_test.rb +++ b/test/openai/internal/util_test.rb @@ -305,6 +305,52 @@ def test_copy_read end end + def test_enum_read_respects_max_len + # rubocop:disable Lint/EmptyBlock + adapter = OpenAI::Internal::Util::ReadIOAdapter.new(%w[abc def].to_enum) {} + # rubocop:enable Lint/EmptyBlock + + assert_equal("", adapter.read(0)) + assert_equal("a", adapter.read(1)) + assert_equal("bcd", adapter.read(3)) + assert_equal("ef", adapter.read(99)) + assert_nil(adapter.read(1)) + end + + def test_enum_read_respects_byte_lengths + input = ["\xC3\xA9b".dup.force_encoding(Encoding::UTF_8)] + # rubocop:disable Lint/EmptyBlock + adapter = OpenAI::Internal::Util::ReadIOAdapter.new(input.to_enum) {} + # rubocop:enable Lint/EmptyBlock + + assert_equal("\xC3".b, adapter.read(1)) + assert_equal("\xA9b".b, adapter.read(2)) + assert_nil(adapter.read(1)) + end + + def test_enum_read_all_includes_buffered_bytes + # rubocop:disable Lint/EmptyBlock + adapter = OpenAI::Internal::Util::ReadIOAdapter.new(%w[abc def].to_enum) {} + # rubocop:enable Lint/EmptyBlock + + assert_equal("ab", adapter.read(2)) + assert_equal("cdef", adapter.read) + assert_equal("", adapter.read) + end + + def test_enum_read_clears_out_string_at_eof + out = +"stale" + # rubocop:disable Lint/EmptyBlock + adapter = OpenAI::Internal::Util::ReadIOAdapter.new(["abc"].to_enum) {} + # rubocop:enable Lint/EmptyBlock + + assert_equal("abc", adapter.read(99, out)) + assert_same(out, adapter.read(0, out)) + assert_equal("", out) + assert_nil(adapter.read(1, out)) + assert_equal("", out) + end + def test_copy_write cases = { StringIO.new => "",