1 files changed, 87 insertions, 5 deletions

diff --git a/paramiko/sftp_client.py b/paramiko/sftp_client.py
index 99a29e3..62127cc 100644
--- a/paramiko/sftp_client.py
+++ b/paramiko/sftp_client.py
@@ -39,6 +39,7 @@ from paramiko.sftp import BaseSFTP, CMD_OPENDIR, CMD_HANDLE, SFTPError, CMD_READ
 from paramiko.sftp_attr import SFTPAttributes
 from paramiko.ssh_exception import SSHException
 from paramiko.sftp_file import SFTPFile
+from paramiko.util import ClosingContextManager
 
 
 def _to_unicode(s):
@@ -58,12 +59,14 @@ def _to_unicode(s):
 b_slash = b'/'
 
 
-class SFTPClient(BaseSFTP):
+class SFTPClient(BaseSFTP, ClosingContextManager):
     """
     SFTP client object.
-    
+
     Used to open an SFTP session across an open SSH `.Transport` and perform
     remote file operations.
+    
+    Instances of this class may be used as context managers.
     """
     def __init__(self, sock):
         """
@@ -98,16 +101,30 @@ class SFTPClient(BaseSFTP):
             raise SSHException('EOF during negotiation')
         self._log(INFO, 'Opened sftp connection (server version %d)' % server_version)
 
-    def from_transport(cls, t):
+    def from_transport(cls, t, window_size=None, max_packet_size=None):
         """
         Create an SFTP client channel from an open `.Transport`.
 
+        Setting the window and packet sizes might affect the transfer speed.
+        The default settings in the `.Transport` class are the same as in
+        OpenSSH and should work adequately for both files transfers and
+        interactive sessions.
+
         :param .Transport t: an open `.Transport` which is already authenticated
+        :param int window_size:
+            optional window size for the `.SFTPClient` session.
+        :param int max_packet_size:
+            optional max packet size for the `.SFTPClient` session..
+
         :return:
             a new `.SFTPClient` object, referring to an sftp session (channel)
             across the transport
+
+        .. versionchanged:: 1.15
+            Added the ``window_size`` and ``max_packet_size`` arguments.
         """
-        chan = t.open_session()
+        chan = t.open_session(window_size=window_size,
+                              max_packet_size=max_packet_size)
         if chan is None:
             return None
         chan.invoke_subsystem('sftp')
@@ -196,6 +213,71 @@ class SFTPClient(BaseSFTP):
         self._request(CMD_CLOSE, handle)
         return filelist
 
+    def listdir_iter(self, path='.', read_aheads=50):
+        """
+        Generator version of `.listdir_attr`.
+
+        See the API docs for `.listdir_attr` for overall details.
+
+        This function adds one more kwarg on top of `.listdir_attr`:
+        ``read_aheads``, an integer controlling how many
+        ``SSH_FXP_READDIR`` requests are made to the server. The default of 50
+        should suffice for most file listings as each request/response cycle
+        may contain multiple files (dependant on server implementation.)
+
+        .. versionadded:: 1.15
+        """
+        path = self._adjust_cwd(path)
+        self._log(DEBUG, 'listdir(%r)' % path)
+        t, msg = self._request(CMD_OPENDIR, path)
+
+        if t != CMD_HANDLE:
+            raise SFTPError('Expected handle')
+
+        handle = msg.get_string()
+
+        nums = list()
+        while True:
+            try:
+                # Send out a bunch of readdir requests so that we can read the
+                # responses later on Section 6.7 of the SSH file transfer RFC
+                # explains this
+                # http://filezilla-project.org/specs/draft-ietf-secsh-filexfer-02.txt
+                for i in range(read_aheads):
+                    num = self._async_request(type(None), CMD_READDIR, handle)
+                    nums.append(num)
+
+
+                # For each of our sent requests
+                # Read and parse the corresponding packets
+                # If we're at the end of our queued requests, then fire off
+                # some more requests
+                # Exit the loop when we've reached the end of the directory
+                # handle
+                for num in nums:
+                    t, pkt_data = self._read_packet()
+                    msg = Message(pkt_data)
+                    new_num = msg.get_int()
+                    if num == new_num:
+                        if t == CMD_STATUS:
+                            self._convert_status(msg)
+                    count = msg.get_int()
+                    for i in range(count):
+                        filename = msg.get_text()
+                        longname = msg.get_text()
+                        attr = SFTPAttributes._from_msg(
+                            msg, filename, longname)
+                        if (filename != '.') and (filename != '..'):
+                            yield attr
+
+                # If we've hit the end of our queued requests, reset nums.
+                nums = list()
+
+            except EOFError:
+                self._request(CMD_CLOSE, handle)
+                return
+
+
     def open(self, filename, mode='r', bufsize=-1):
         """
         Open a file on the remote server.  The arguments are the same as for
@@ -578,7 +660,7 @@ class SFTPClient(BaseSFTP):
 
         .. versionadded:: 1.4
         .. versionchanged:: 1.7.4
-            ``callback`` and rich attribute return value added.   
+            ``callback`` and rich attribute return value added.
         .. versionchanged:: 1.7.7
             ``confirm`` param added.
         """