Merge pull request #81 from jstuckey/add-limit-argument

Make file operation work with CSV library

Author: mfazekas

lib/net/sftp/operations/file.rb

@@ -81,21 +81,35 @@ def read(n=nil)
     # Reads up to the next instance of +sep_string+ in the stream, and
     # returns the bytes read (including +sep_string+). If +sep_string+ is
     # omitted, it defaults to +$/+. If EOF is encountered before any data
-    # could be read, #gets will return +nil+.
-    def gets(sep_string=$/)
-      delim = if sep_string.length == 0
+    # could be read, #gets will return +nil+. If the first argument is an
+    # integer, or optional second argument is given, the returning string
+    # would not be longer than the given value in bytes.
+    def gets(sep_or_limit=$/, limit=Float::INFINITY)
+      if sep_or_limit.is_a? Integer
+        sep_string = $/
+        lim = sep_or_limit
+      else
+        sep_string = sep_or_limit
+        lim = limit
+      end
+
+      delim = if sep_string && sep_string.length == 0
         "#{$/}#{$/}"
       else
         sep_string
       end
 
       loop do
-        at = @buffer.index(delim)
+        at = @buffer.index(delim) if delim
         if at
-          offset = at + delim.length
+          offset = [at + delim.length, lim].min
           @pos += offset
           line, @buffer = @buffer[0,offset], @buffer[offset..-1]
           return line
+        elsif lim < @buffer.length
+          @pos += lim
+          line, @buffer = @buffer[0,lim], @buffer[lim..-1]
+          return line
         elsif !fill
           return nil if @buffer.empty?
           @pos += @buffer.length
@@ -107,12 +121,20 @@ def gets(sep_string=$/)
 
     # Same as #gets, but raises EOFError if EOF is encountered before any
     # data could be read.
-    def readline(sep_string=$/)
-      line = gets(sep_string)
+    def readline(sep_or_limit=$/, limit=Float::INFINITY)
+      line = gets(sep_or_limit, limit)
       raise EOFError if line.nil?
       return line
     end
 
+    # Resets position to beginning of file
+    def rewind
+      @pos      = 0
+      @real_pos = 0
+      @real_eof = false
+      @buffer   = ""
+    end
+
     # Writes the given data to the stream, incrementing the file position and
     # returning the number of bytes written.
     def write(data)

test/test_file.rb

@@ -91,6 +91,37 @@ def test_gets_when_no_such_delimiter_exists_in_stream_should_read_to_EOF
     assert @file.eof?
   end
 
+  def test_gets_when_nil_delimiter_should_fread_to_EOF
+    @sftp.expects(:read!).times(2).returns("hello world\ngoodbye world\n\nfarewell!\n", nil)
+    assert_equal "hello world\ngoodbye world\n\nfarewell!\n", @file.gets(nil)
+    assert @file.eof?
+  end
+
+  def test_gets_with_integer_argument_should_read_number_of_bytes
+    @sftp.expects(:read!).returns("hello world\ngoodbye world\n\nfarewell!\n")
+    assert_equal "hello w", @file.gets(7)
+  end
+
+  def test_gets_with_delimiter_and_limit_should_read_to_delimiter_if_less_than_limit
+    @sftp.expects(:read!).returns("hello world\ngoodbye world\n\nfarewell!\n")
+    assert_equal "hello w", @file.gets("w", 11)
+  end
+
+  def test_gets_with_delimiter_and_limit_should_read_to_limit_if_less_than_delimiter
+    @sftp.expects(:read!).returns("hello world\ngoodbye world\n\nfarewell!\n")
+    assert_equal "hello", @file.gets("w", 5)
+  end
+
+  def test_gets_when_no_such_delimiter_exists_in_stream_but_limit_provided_should_read_to_limit
+    @sftp.expects(:read!).returns("hello world\ngoodbye world\n\nfarewell!\n")
+    assert_equal "hello w", @file.gets("z", 7)
+  end
+
+  def test_gets_when_nil_delimiter_and_limit_provided_should_read_to_limit
+    @sftp.expects(:read!).returns("hello world\ngoodbye world\n\nfarewell!\n")
+    assert_equal "hello w", @file.gets(nil, 7)
+  end
+
   def test_gets_at_EOF_should_return_nil
     @sftp.expects(:read!).returns(nil)
     assert_nil @file.gets
@@ -102,6 +133,15 @@ def test_readline_should_raise_exception_on_EOF
     assert_raises(EOFError) { @file.readline }
   end
 
+  def test_rewind_should_reset_to_beginning_of_file
+    @sftp.expects(:read!).times(2).returns("hello world", nil)
+    @file.read
+    assert @file.eof?
+    @file.rewind
+    assert !@file.eof?
+    assert_equal 0, @file.pos
+  end
+
   def test_write_should_write_data_and_increment_pos_and_return_data_length
     @sftp.expects(:write!).with("handle", 0, "hello world")
     assert_equal 11, @file.write("hello world")