Add limit argument to gets method

Ruby's IO class defines the #gets method as having an optional limit
argument. This can either be passed in as an integer for the first
argument or passed in addition to the delimiter. This makes argument
handling a little awkward unfortunately.

Author: jstuckey

lib/net/sftp/operations/file.rb

@@ -82,7 +82,15 @@ def read(n=nil)
     # 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=$/)
+    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.length == 0
         "#{$/}#{$/}"
       else
@@ -92,7 +100,7 @@ def gets(sep_string=$/)
       loop do
         at = @buffer.index(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

test/test_file.rb

@@ -91,6 +91,21 @@ def test_gets_when_no_such_delimiter_exists_in_stream_should_read_to_EOF
     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_at_EOF_should_return_nil
     @sftp.expects(:read!).returns(nil)
     assert_nil @file.gets
@@ -156,4 +171,4 @@ def test_stat_should_return_attributes_object_for_handle
     @sftp.expects(:fstat!).with("handle").returns(stat)
     assert_equal stat, @file.stat
   end
-end
+end