<commit_before># Stream of objects, with any underlying IO: File, Pipe, Socket, StringIO.
# Stream is bidirectional if the IO is bidirectional. Serializes objects using
# any of several serializers: marshal, yaml, json, msgpack. Works with
# select/readpartial if the serializer supports it (msgpack and yajl do).
# ObjectStream supports three styles of iteration: Enumerable, blocking read,
# and yielding read. The #expect method is a way to instantiate custom classes
# in cases (msgpack, json) which do not support it natively. The #consume
# method helps with handshake protocols.
module ObjectStream
  include Enumerable
  
  attr_reader :io
  attr_accessor :peer_name
  
  MARSHAL_TYPE  = "marshal"
  YAML_TYPE     = "yaml"
  JSON_TYPE     = "json"
  MSGPACK_TYPE  = "msgpack"
  
  TYPES = [
    MARSHAL_TYPE, YAML_TYPE, JSON_TYPE, MSGPACK_TYPE
  ]

  MAX_OUTBOX = 10

  # Raised when maxbuf exceeded.
  class OverflowError < StandardError; end

  @stream_class_map =
    Hash.new {|h,type| raise ArgumentError, "unknown type: #{type.inspect}"}
  @mutex = Mutex.new

  class << self
    def new io, type: MARSHAL_TYPE, **opts
      if io.kind_of? ObjectStream
        raise ArgumentError,
          "given io is already an ObjectStream: #{io.inspect}"
      end
      ## discover type by reading the io
      ## can the io be an in-memory queue?
      stream_class_for(type).new io, **opts
    end

    def stream_class_for type
      cl = @stream_class_map[type]
      if cl.respond_to? :new
        cl
      else
        @mutex.synchronize do ## seems like overkill. Is autoload working?
          if cl.respond_to? :new
            cl
          else
            @stream_class_map[type] = cl.call
          end
        end
      end
    end
    
    def register_type type, &bl
      @stream_class_map[type] = bl
    end
  end
  
  def initialize io, **opts
    @io = io
    @inbox = nil
    @outbox = []
    @peer_name = "unknown"
    @consumers = []
    unexpect
  end
  
  def to_s
    "#<#{self.class} to #{peer_name}, io=#{io.inspect}>"
  end
  
  def expect cl = nil; end
  def unexpect; expect nil; end

  def consume &bl
    @consumers << bl
  end

  def try_consume obj
    if bl = @consumers.shift
      bl[obj]
      true
    else
      false
    end
  end

  # raises EOFError
  def read
    if block_given?
      read_from_inbox {|obj| try_consume(obj) or yield obj}
      read_from_stream {|obj| try_consume(obj) or yield obj}
      return nil
    else
      read_one
    end
  end

  # read and return exactly one; blocking
  def read_one
    if @inbox and not @inbox.empty?
      return @inbox.shift
    end
    
    have_result = false
    result = nil
    until have_result
      read do |obj| # might not read enough bytes to yield an obj
        if have_result
          (@inbox||=[]) << obj
        else
          have_result = true
          result = obj
        end
      end
    end
    result
  end

  def read_from_inbox
    if @inbox and not @inbox.empty?
      @inbox.each {|obj| yield obj}
      @inbox.clear
    end
  end
  
  def write object
    write_to_buffer object
    flush_buffer
  end
  alias << write

  def write_to_outbox object=nil, &bl
    @outbox << (bl || object)
    flush_outbox if @outbox.size > MAX_OUTBOX
    self
  end

  def flush_outbox
    @outbox.each do |object|
      object = object.call if object.kind_of? Proc
      write_to_stream object
    end
    @outbox.clear
    self
  end
<commit_msg>max_outbox is settable per stream<commit_after># Stream of objects, with any underlying IO: File, Pipe, Socket, StringIO.
# Stream is bidirectional if the IO is bidirectional. Serializes objects using
# any of several serializers: marshal, yaml, json, msgpack. Works with
# select/readpartial if the serializer supports it (msgpack and yajl do).
# ObjectStream supports three styles of iteration: Enumerable, blocking read,
# and yielding read. The #expect method is a way to instantiate custom classes
# in cases (msgpack, json) which do not support it natively. The #consume
# method helps with handshake protocols.
module ObjectStream
  include Enumerable
  
  attr_reader :io
  attr_reader :max_outbox
  attr_accessor :peer_name
  
  MARSHAL_TYPE  = "marshal"
  YAML_TYPE     = "yaml"
  JSON_TYPE     = "json"
  MSGPACK_TYPE  = "msgpack"
  
  TYPES = [
    MARSHAL_TYPE, YAML_TYPE, JSON_TYPE, MSGPACK_TYPE
  ]

  DEFAULT_MAX_OUTBOX = 10

  # Raised when maxbuf exceeded.
  class OverflowError < StandardError; end

  @stream_class_map =
    Hash.new {|h,type| raise ArgumentError, "unknown type: #{type.inspect}"}
  @mutex = Mutex.new

  class << self
    def new io, type: MARSHAL_TYPE, **opts
      if io.kind_of? ObjectStream
        raise ArgumentError,
          "given io is already an ObjectStream: #{io.inspect}"
      end
      ## discover type by reading the io
      ## can the io be an in-memory queue?
      stream_class_for(type).new io, **opts
    end

    def stream_class_for type
      cl = @stream_class_map[type]
      if cl.respond_to? :new
        cl
      else
        @mutex.synchronize do ## seems like overkill. Is autoload working?
          if cl.respond_to? :new
            cl
          else
            @stream_class_map[type] = cl.call
          end
        end
      end
    end
    
    def register_type type, &bl
      @stream_class_map[type] = bl
    end
  end
  
  def initialize io, max_outbox: DEFAULT_MAX_OUTBOX, **opts
    @io = io
    @max_outbox = max_outbox
    @inbox = nil
    @outbox = []
    @peer_name = "unknown"
    @consumers = []
    unexpect
  end
  
  def to_s
    "#<#{self.class} to #{peer_name}, io=#{io.inspect}>"
  end
  
  def expect cl = nil; end
  def unexpect; expect nil; end

  def consume &bl
    @consumers << bl
  end

  def try_consume obj
    if bl = @consumers.shift
      bl[obj]
      true
    else
      false
    end
  end

  # raises EOFError
  def read
    if block_given?
      read_from_inbox {|obj| try_consume(obj) or yield obj}
      read_from_stream {|obj| try_consume(obj) or yield obj}
      return nil
    else
      read_one
    end
  end

  # read and return exactly one; blocking
  def read_one
    if @inbox and not @inbox.empty?
      return @inbox.shift
    end
    
    have_result = false
    result = nil
    until have_result
      read do |obj| # might not read enough bytes to yield an obj
        if have_result
          (@inbox||=[]) << obj
        else
          have_result = true
          result = obj
        end
      end
    end
    result
  end

  def read_from_inbox
    if @inbox and not @inbox.empty?
      @inbox.each {|obj| yield obj}
      @inbox.clear
    end
  end
  
  def write object
    write_to_buffer object
    flush_buffer
  end
  alias << write

  def write_to_outbox object=nil, &bl
    @outbox << (bl || object)
    flush_outbox if @outbox.size > max_outbox
    self
  end

  def flush_outbox
    @outbox.each do |object|
      object = object.call if object.kind_of? Proc
      write_to_stream object
    end
    @outbox.clear
    self
  end

{'fraction_non_alphanumeric': 0.07572098309418951, 'fraction_numerical': 0.000568262537292229, 'mean_word_length': 2.346007604562738, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '617690', 'n_tokens_mistral': 2359, 'n_tokens_neox': 2294, 'n_words': 752}