traffic learner: avoid truncated output files on error

add_argument(type=argparse.FileType('w'), ...) will open the file
and leave it empty if the script fails.

Signed-off-by: Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>

1 files changed, 8 insertions, 1 deletions

diff --git a/script/traffic_learner b/script/traffic_learner
index b46521f70cd..cad1a64a133 100755
--- a/script/traffic_learner
+++ b/script/traffic_learner
@@ -31,7 +31,7 @@ info = logger.info
 def main():
     parser = argparse.ArgumentParser(description=__doc__,
                         formatter_class=argparse.RawDescriptionHelpFormatter)
-    parser.add_argument('-o', '--out', type=argparse.FileType('w'),
+    parser.add_argument('-o', '--out',
                         help="write model here")
     parser.add_argument('--dns-mode', choices=['inline', 'count'],
                         help='how to deal with DNS', default='count')
@@ -45,6 +45,13 @@ def main():
         error("Please specify a filename using the --out option.")
         return 1
 
+    try:
+        outfile = open(args.out, 'w')
+    except IOError as e:
+        error("could not open %s" % args.out)
+        error(e)
+        return 1
+
     if args.SUMMARY_FILE is sys.stdin:
         info("reading from STDIN...")