diff --git a/tftpy/TftpClient.py b/tftpy/TftpClient.py index 9ea905f..f84d801 100644 --- a/tftpy/TftpClient.py +++ b/tftpy/TftpClient.py @@ -27,15 +27,15 @@ class TftpClient(TftpSession): def download(self, filename, output, packethook=None, timeout=SOCK_TIMEOUT): """This method initiates a tftp download from the configured remote - host, requesting the filename passed. It saves the file to a local - file specified in the output parameter. If a packethook is provided, - it must be a function that takes a single parameter, which will be a - copy of each DAT packet received in the form of a TftpPacketDAT - object. The timeout parameter may be used to override the default - SOCK_TIMEOUT setting, which is the amount of time that the client will - wait for a receive packet to arrive. + host, requesting the filename passed. It writes the file to output, + which can be a file-like object or a path to a local file. If a + packethook is provided, it must be a function that takes a single + parameter, which will be a copy of each DAT packet received in the + form of a TftpPacketDAT object. The timeout parameter may be used to + override the default SOCK_TIMEOUT setting, which is the amount of time + that the client will wait for a receive packet to arrive. - Note: If output is a hyphen then stdout is used.""" + Note: If output is a hyphen, stdout is used.""" # We're downloading. log.debug("Creating download context with the following params:") log.debug("host = %s, port = %s, filename = %s, output = %s" @@ -67,17 +67,15 @@ class TftpClient(TftpSession): def upload(self, filename, input, packethook=None, timeout=SOCK_TIMEOUT): """This method initiates a tftp upload to the configured remote host, - uploading the filename passed. If a packethook is provided, it must - be a function that takes a single parameter, which will be a copy of - each DAT packet sent in the form of a TftpPacketDAT object. The - timeout parameter may be used to override the default SOCK_TIMEOUT - setting, which is the amount of time that the client will wait for a - DAT packet to be ACKd by the server. + uploading the filename passed. It reads the file from input, which + can be a file-like object or a path to a local file. If a packethook + is provided, it must be a function that takes a single parameter, + which will be a copy of each DAT packet sent in the form of a + TftpPacketDAT object. The timeout parameter may be used to override + the default SOCK_TIMEOUT setting, which is the amount of time that + the client will wait for a DAT packet to be ACKd by the server. - The input option is the full path to the file to upload, which can - optionally be '-' to read from stdin. - - Note: If output is a hyphen then stdout is used.""" + Note: If input is a hyphen, stdin is used.""" self.context = TftpContextClientUpload(self.host, self.iport, filename, diff --git a/tftpy/TftpContexts.py b/tftpy/TftpContexts.py index c3a1bd4..7f118d8 100644 --- a/tftpy/TftpContexts.py +++ b/tftpy/TftpContexts.py @@ -255,7 +255,11 @@ class TftpContextClientUpload(TftpContext): self.file_to_transfer = filename self.options = options self.packethook = packethook - if input == '-': + # If the input object has a read() function, + # assume it is file-like. + if hasattr(input, 'read'): + self.fileobj = input + elif input == '-': self.fileobj = sys.stdin else: self.fileobj = open(input, "rb") @@ -328,10 +332,12 @@ class TftpContextClientDownload(TftpContext): self.file_to_transfer = filename self.options = options self.packethook = packethook - # FIXME - need to support alternate return formats than files? - # File-like objects would be ideal, ala duck-typing. - # If the filename is -, then use stdout - if output == '-': + # If the output object has a write() function, + # assume it is file-like. + if hasattr(output, 'write'): + self.fileobj = output + # If the output filename is -, then use stdout + elif output == '-': self.fileobj = sys.stdout else: self.fileobj = open(output, "wb")