Allow file objects to be used for download output or upload input.
Objects are considered file-like if they have read() or write() functions. If they do not, they are assumed to be strings to a local path (existing behavior).
This commit is contained in:
parent
05e56b66a4
commit
39b877258b
2 changed files with 27 additions and 23 deletions
|
@ -27,15 +27,15 @@ class TftpClient(TftpSession):
|
||||||
|
|
||||||
def download(self, filename, output, packethook=None, timeout=SOCK_TIMEOUT):
|
def download(self, filename, output, packethook=None, timeout=SOCK_TIMEOUT):
|
||||||
"""This method initiates a tftp download from the configured remote
|
"""This method initiates a tftp download from the configured remote
|
||||||
host, requesting the filename passed. It saves the file to a local
|
host, requesting the filename passed. It writes the file to output,
|
||||||
file specified in the output parameter. If a packethook is provided,
|
which can be a file-like object or a path to a local file. If a
|
||||||
it must be a function that takes a single parameter, which will be a
|
packethook is provided, it must be a function that takes a single
|
||||||
copy of each DAT packet received in the form of a TftpPacketDAT
|
parameter, which will be a copy of each DAT packet received in the
|
||||||
object. The timeout parameter may be used to override the default
|
form of a TftpPacketDAT object. The timeout parameter may be used to
|
||||||
SOCK_TIMEOUT setting, which is the amount of time that the client will
|
override the default SOCK_TIMEOUT setting, which is the amount of time
|
||||||
wait for a receive packet to arrive.
|
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.
|
# We're downloading.
|
||||||
log.debug("Creating download context with the following params:")
|
log.debug("Creating download context with the following params:")
|
||||||
log.debug("host = %s, port = %s, filename = %s, output = %s"
|
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):
|
def upload(self, filename, input, packethook=None, timeout=SOCK_TIMEOUT):
|
||||||
"""This method initiates a tftp upload to the configured remote host,
|
"""This method initiates a tftp upload to the configured remote host,
|
||||||
uploading the filename passed. If a packethook is provided, it must
|
uploading the filename passed. It reads the file from input, which
|
||||||
be a function that takes a single parameter, which will be a copy of
|
can be a file-like object or a path to a local file. If a packethook
|
||||||
each DAT packet sent in the form of a TftpPacketDAT object. The
|
is provided, it must be a function that takes a single parameter,
|
||||||
timeout parameter may be used to override the default SOCK_TIMEOUT
|
which will be a copy of each DAT packet sent in the form of a
|
||||||
setting, which is the amount of time that the client will wait for a
|
TftpPacketDAT object. The timeout parameter may be used to override
|
||||||
DAT packet to be ACKd by the server.
|
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
|
Note: If input is a hyphen, stdin is used."""
|
||||||
optionally be '-' to read from stdin.
|
|
||||||
|
|
||||||
Note: If output is a hyphen then stdout is used."""
|
|
||||||
self.context = TftpContextClientUpload(self.host,
|
self.context = TftpContextClientUpload(self.host,
|
||||||
self.iport,
|
self.iport,
|
||||||
filename,
|
filename,
|
||||||
|
|
|
@ -255,7 +255,11 @@ class TftpContextClientUpload(TftpContext):
|
||||||
self.file_to_transfer = filename
|
self.file_to_transfer = filename
|
||||||
self.options = options
|
self.options = options
|
||||||
self.packethook = packethook
|
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
|
self.fileobj = sys.stdin
|
||||||
else:
|
else:
|
||||||
self.fileobj = open(input, "rb")
|
self.fileobj = open(input, "rb")
|
||||||
|
@ -328,10 +332,12 @@ class TftpContextClientDownload(TftpContext):
|
||||||
self.file_to_transfer = filename
|
self.file_to_transfer = filename
|
||||||
self.options = options
|
self.options = options
|
||||||
self.packethook = packethook
|
self.packethook = packethook
|
||||||
# FIXME - need to support alternate return formats than files?
|
# If the output object has a write() function,
|
||||||
# File-like objects would be ideal, ala duck-typing.
|
# assume it is file-like.
|
||||||
# If the filename is -, then use stdout
|
if hasattr(output, 'write'):
|
||||||
if output == '-':
|
self.fileobj = output
|
||||||
|
# If the output filename is -, then use stdout
|
||||||
|
elif output == '-':
|
||||||
self.fileobj = sys.stdout
|
self.fileobj = sys.stdout
|
||||||
else:
|
else:
|
||||||
self.fileobj = open(output, "wb")
|
self.fileobj = open(output, "wb")
|
||||||
|
|
Reference in a new issue