3rdparty/modules/aviator/lib/puppet/feature/composite_io.rb

   1 #--
   2 # Copyright (c) 2007-2012 Nick Sieger.
   3 # See the file README.txt included with the distribution for
   4 # software license details.
   5 #++
   6
   7 # Concatenate together multiple IO objects into a single, composite IO object
   8 # for purposes of reading as a single stream.
   9 #
  10 # Usage:
  11 #
  12 #     crio = CompositeReadIO.new(StringIO.new('one'), StringIO.new('two'), StringIO.new('three'))
  13 #     puts crio.read # => "onetwothree"
  14 #
  15 class CompositeReadIO
  16   # Create a new composite-read IO from the arguments, all of which should
  17   # respond to #read in a manner consistent with IO.
  18   def initialize(*ios)
  19     @ios = ios.flatten
  20     @index = 0
  21   end
  22
  23   # Read from IOs in order until `length` bytes have been received.
  24   def read(length = nil, outbuf = nil)
  25     got_result = false
  26     outbuf = outbuf ? outbuf.replace("") : ""
  27
  28     while io = current_io
  29       if result = io.read(length)
  30         got_result ||= !result.nil?
  31         result.force_encoding("BINARY") if result.respond_to?(:force_encoding)
  32         outbuf << result
  33         length -= result.length if length
  34         break if length == 0
  35       end
  36       advance_io
  37     end
  38     (!got_result && length) ? nil : outbuf
  39   end
  40
  41   def rewind
  42     @ios.each { |io| io.rewind }
  43     @index = 0
  44   end
  45
  46   private
  47
  48   def current_io
  49     @ios[@index]
  50   end
  51
  52   def advance_io
  53     @index += 1
  54   end
  55 end
  56
  57 # Convenience methods for dealing with files and IO that are to be uploaded.
  58 class UploadIO
  59   # Create an upload IO suitable for including in the params hash of a
  60   # Net::HTTP::Post::Multipart.
  61   #
  62   # Can take two forms. The first accepts a filename and content type, and
  63   # opens the file for reading (to be closed by finalizer).
  64   #
  65   # The second accepts an already-open IO, but also requires a third argument,
  66   # the filename from which it was opened (particularly useful/recommended if
  67   # uploading directly from a form in a framework, which often save the file to
  68   # an arbitrarily named RackMultipart file in /tmp).
  69   #
  70   # Usage:
  71   #
  72   #     UploadIO.new("file.txt", "text/plain")
  73   #     UploadIO.new(file_io, "text/plain", "file.txt")
  74   #
  75   attr_reader :content_type, :original_filename, :local_path, :io, :opts
  76
  77   def initialize(filename_or_io, content_type, filename = nil, opts = {})
  78     io = filename_or_io
  79     local_path = ""
  80     if io.respond_to? :read
  81       # in Ruby 1.9.2, StringIOs no longer respond to path
  82       # (since they respond to :length, so we don't need their local path, see parts.rb:41)
  83       local_path = filename_or_io.respond_to?(:path) ? filename_or_io.path : "local.path"
  84     else
  85       io = File.open(filename_or_io)
  86       local_path = filename_or_io
  87     end
  88     filename ||= local_path
  89
  90     @content_type = content_type
  91     @original_filename = File.basename(filename)
  92     @local_path = local_path
  93     @io = io
  94     @opts = opts
  95   end
  96
  97   def self.convert!(io, content_type, original_filename, local_path)
  98     raise ArgumentError, "convert! has been removed. You must now wrap IOs using:\nUploadIO.new(filename_or_io, content_type, filename=nil)\nPlease update your code."
  99   end
 100
 101   def method_missing(*args)
 102     @io.send(*args)
 103   end
 104
 105   def respond_to?(meth, include_all = false)
 106     @io.respond_to?(meth, include_all) || super(meth, include_all)
 107   end
 108 end