Take advantage of the automatic file handling thinfy in argparser
This commit is contained in:
parent
e6b3345421
commit
f390da248c
1 changed files with 6 additions and 11 deletions
|
@ -34,10 +34,13 @@ def main():
|
|||
|
||||
# -o can be used to write the .sshwot thing to a file instead of stdout
|
||||
parser.add_argument('-o',
|
||||
# We store one argument given after this one to the property
|
||||
# outfile, which will be None otherwise
|
||||
# Given one argument, we open a file of that name and store it
|
||||
# to outfile, which will be sys.stdout.buffer otherwise.
|
||||
# We use .buffer since we're going to write binary data
|
||||
action = 'store',
|
||||
dest = 'outfile',
|
||||
type = argparse.FileType('wb'),
|
||||
default = sys.stdout.buffer,
|
||||
# This is what will be displayed in the help after -o
|
||||
metavar = 'outfile',
|
||||
help = 'write the sshwot file to a given file instead of the stdout'
|
||||
|
@ -68,15 +71,7 @@ def main():
|
|||
for known_hosts_entry in known_host_entries:
|
||||
entries.append(process_known_hosts.known_hosts_to_entry(known_hosts_entry))
|
||||
|
||||
if args.outfile is None:
|
||||
# Write to stdout by default
|
||||
# We use sys.stdout.buffer instead of just sys.stdout because we
|
||||
# are writing binary data
|
||||
write_file.write(sys.stdout.buffer, entries)
|
||||
|
||||
else:
|
||||
with open(args.outfile, 'wb') as outf:
|
||||
write_file.write(outf, entries)
|
||||
write_file.write(args.outfile, entries)
|
||||
|
||||
if __name__ == '__main__':
|
||||
try:
|
||||
|
|
Loading…
Reference in a new issue