lib/bencode/parser.rb in bencode-0.8.0 vs lib/bencode/parser.rb in bencode-0.8.1

@@ -1,24 +1,24 @@
 require 'stringio'
-require 'bencode/parseio'
 
 module BEncode
   class Parser
     attr_reader :stream
 
     def initialize(stream)
-      if stream.kind_of? BEncode::ParseIO
-        @stream = stream
-      elsif stream.respond_to? :string
-        @stream = StringIO.new stream.string
-      elsif stream.respond_to? :to_s
-        @stream = StringIO.new stream.to_s
-      end
+      @stream =
+        if stream.kind_of?(IO) || stream.kind_of?(StringIO)
+          stream
+        elsif stream.respond_to? :string
+          StringIO.new stream.string
+        elsif stream.respond_to? :to_s
+          StringIO.new stream.to_s
+        end
     end
 
     def parse!
-      case stream.peek
+      case peek
         when ?i then parse_integer!
         when ?l then parse_list!
         when ?d then parse_dict!
         when ?0 .. ?9 then parse_string!
       end
@@ -37,19 +37,19 @@
     end
 
     def parse_list!
       stream.getc
       ary = []
-      ary.push(parse!) until stream.peek == ?e
+      ary.push(parse!) until peek == ?e
       stream.getc
       ary
     end
 
     def parse_dict!
       stream.getc
       hsh = {}
-      until stream.peek == ?e
+      until peek == ?e
         key = parse!
 
         unless key.is_a? String or key.is_a? Fixnum
           raise BEncode::DecodeError, "key must be a string or number"
         end
@@ -73,8 +73,14 @@
       rescue
         raise BEncode::DecodeError, "invalid string length"
       end
 
       str
+    end
+    
+    def peek
+      c = stream.getc
+      stream.ungetc(c)
+      c
     end
   end
 end