python open_sftp() context manager for sftp read and writing of files with similar behavior to open()

open_sftp.py

from contextlib import contextmanager
from urlparse import urlparse

from paramiko import AuthenticationException
import pysftp


@contextmanager
def open_sftp(url, mode='r', n_retries=5, retry_sleep_time=5):
    """Context manager to read/write a file via SFTP

    Note: We are using urlparse.urlparse to parse the :attr:`url`. There is
    an corner case where this will not work if there is a ':' in the username.

    :param url: Path to file with the format
                sftp://username:password@host:[port]/path/to/file.txt
    :param mode: The mode to open the file with. Either 'r' or 'w'
    :param n_retries: The number of times to retry connecting to the SFTP server
    :param retry_sleep_time: The number of seconds to sleep between retries
    :returns: File object for the specified file
    :raises:
        :ValueError: If there is an error with the input
        :AuthenticationException: If could not connect
        :IOError: If file cannot be found/opened
    """
    # Make sure that we have a proper minimum retry settings
    n_retries = n_retries or 5
    n_retries = max(n_retries, 1)
    retry_sleep_time = retry_sleep_time or 5
    retry_sleep_time = max(retry_sleep_time, 1)

    # Parse url into username, password, host, port, path
    # using the proper ftp URL syntax of
    # sftp://username:password@host:port/path/to/file.txt
    split_url = urlparse(url)

    if not split_url.scheme == 'sftp':
        raise ValueError(
            "Invalid URL: Invalid scheme '{}'".format(split_url.scheme))
    if not split_url.hostname:
        raise ValueError("Invalid URL: No hostname")
    if not split_url.path:
        raise ValueError("Invalid URL: No file path")

    path_split = split_url.path.split('/')
    if len(path_split) < 1:
        raise ValueError(
            "Invalid URL: Invalid path '{}'".format(split_url.path))
    file_name = path_split.pop()
    file_path = path_split
    if not file_name:
        raise ValueError("Invalid URL: No file name")

    cinfo = {
        'host': split_url.hostname,
        'username': split_url.username,
        'password': split_url.password
    }
    try:
        if split_url.port:
            cinfo['port'] = split_url.port
    except ValueError:
        raise ValueError("Invalid URL: Invalid port")

    # Connect to host. Sometimes it takes a couple tries... =(
    sftp = None
    for i in xrange(n_retries-1):
        try:
            sftp = pysftp.Connection(**cinfo)
            break
        except AuthenticationException:
            warnings.warn("Could not authenticate... Will re-attempt")
            import time as t
            t.sleep(retry_sleep_time)

    # Try one more time... if this fails, raise the AuthenticationError
    if not sftp:
        sftp = pysftp.Connection(**cinfo)

    try:
        for directory in file_path:

            # Skip blanks (e.g. caused by '//')
            if not directory:
                continue

            try:
                sftp.chdir(directory)
            except IOError as exc:
                raise IOError(
                    "Could not navigate to path '{}': {}".format(
                        '/'.join(file_path), exc))

        try:
            file_obj = sftp.open(file_name, mode=mode)
        except IOError as exc:
            raise IOError("Could not open file '{}' in path '{}': {}".format(
                file_name, file_path, exc))

        yield file_obj
    finally:
        sftp.close()