[tor-commits] [sbws/master] Move back write function to V3BWFile

[pastly at torproject.org]

commit 666b9d7f57f2080146c9ee14aa8ee59e9228849b
Author: juga0 <juga at riseup.net>
Date:   Mon Jun 25 15:52:16 2018 +0000

    Move back write function to V3BWFile
    
    but not allowing /dev/stdout
---
 sbws/core/generate.py | 28 +---------------------------
 sbws/lib/v3bwfile.py  | 23 ++++++++++++++++++++++-
 2 files changed, 23 insertions(+), 28 deletions(-)

diff --git a/sbws/core/generate.py b/sbws/core/generate.py
index 827058e..6492423 100644
--- a/sbws/core/generate.py
+++ b/sbws/core/generate.py
@@ -4,7 +4,6 @@ from sbws.lib.resultdump import load_recent_results_in_datadir
 from argparse import ArgumentDefaultsHelpFormatter
 import os
 import logging
-from sbws.util.filelock import DirectoryLock
 from sbws.util.timestamp import now_fname
 
 log = logging.getLogger(__name__)
@@ -33,31 +32,6 @@ def gen_parser(sub):
                    'out, and we do so proportionally')
 
 
-def _write_v3bw_file(v3bwfile, output):
-    log.info('Writing v3bw file to %s', output)
-    out_dir = os.path.dirname(output)
-    out_link = os.path.join(out_dir, 'latest.v3bw')
-    if os.path.exists(out_link):
-        log.debug('Deleting existing symlink before creating a new one.')
-        os.remove(out_link)
-    # to keep test_generate.py working
-    if output != '/dev/stdout':
-        with DirectoryLock(out_dir):
-            with open(output, 'wt') as fd:
-                fd.write(str(v3bwfile.header))
-                for line in v3bwfile.bw_lines:
-                    fd.write(str(line))
-            output_basename = os.path.basename(output)
-            log.debug('Creating symlink from {} to {}.'
-                      .format(output_basename, out_link))
-            os.symlink(output_basename, out_link)
-    else:
-        with open(output, 'wt') as fd:
-            fd.write(str(v3bwfile.header))
-            for line in v3bwfile.bw_lines:
-                fd.write(str(line))
-
-
 def main(args, conf):
     if not is_initted(args.directory):
         fail_hard('Sbws isn\'t initialized.  Try sbws init')
@@ -83,5 +57,5 @@ def main(args, conf):
         return
     bw_file = V3BWFile.from_arg_results(args, conf, results)
     output = args.output or conf['paths']['v3bw_fname'].format(now_fname())
-    _write_v3bw_file(bw_file, output)
+    V3BWFile.write(output)
     log.info('Mean bandwidth per line: %f "KiB"', bw_file.avg_bw)
diff --git a/sbws/lib/v3bwfile.py b/sbws/lib/v3bwfile.py
index 64395ef..fb9e3d3 100644
--- a/sbws/lib/v3bwfile.py
+++ b/sbws/lib/v3bwfile.py
@@ -3,12 +3,13 @@
 (v3bw) used by bandwidth authorities."""
 
 import logging
+import os
 from statistics import median
 
 from sbws import __version__
 from sbws.globals import SPEC_VERSION, BW_LINE_SIZE
 from sbws.lib.resultdump import ResultSuccess, _ResultType
-from sbws.util.filelock import FileLock
+from sbws.util.filelock import FileLock, DirectoryLock
 from sbws.util.timestamp import now_isodt_str, unixts_to_isodt_str
 
 log = logging.getLogger(__name__)
@@ -405,3 +406,23 @@ class V3BWFile(object):
         header = V3BWHeader.from_results(conf, results)
         f = cls(header, bw_lines)
         return f
+
+    def write(self, output):
+        if output == '/dev/stdout':
+            log.info("Writing to stdout is not supported.")
+            return
+        log.info('Writing v3bw file to %s', output)
+        out_dir = os.path.dirname(output)
+        out_link = os.path.join(out_dir, 'latest.v3bw')
+        if os.path.exists(out_link):
+            log.debug('Deleting existing symlink before creating a new one.')
+            os.remove(out_link)
+        with DirectoryLock(out_dir):
+            with open(output, 'wt') as fd:
+                fd.write(str(self.header))
+                for line in self.bw_lines:
+                    fd.write(str(line))
+            output_basename = os.path.basename(output)
+            log.debug('Creating symlink from {} to {}.'
+                      .format(output_basename, out_link))
+            os.symlink(output_basename, out_link)