switch to argparse

1 files changed, 104 insertions, 124 deletions

diff --git a/bin/msp430-etv b/bin/msp430-etv
index c3a0f5e..18d7622 100755
--- a/bin/msp430-etv
+++ b/bin/msp430-etv
@@ -1,29 +1,7 @@
 #!/usr/bin/env python3
 # vim:tabstop=4:softtabstop=4:shiftwidth=4:textwidth=160:smarttab:expandtab
 
-import getopt
-import itertools
-import matplotlib.pyplot as plt
-import numpy as np
-import os
-import re
-from shutil import which
-import subprocess
-import sys
-import tempfile
-import time
-
-opt = dict()
-
-
-def show_help():
-    print(
-        """msp430-etv - MSP430 EnergyTrace Visualizer
-
-USAGE
-
-msp430-etv [--load <file> | <measurement duration>] [--save <file>]
-    [--skip <count>] [--threshold <power>] [--plot=U|I|P] [--stat]
+"""msp430-etv - MSP430 EnergyTrace Visualizer
 
 DESCRIPTION
 
@@ -34,47 +12,28 @@ specifying <measurement duration> in seconds) or loaded from a logfile using
 
 This program is not affiliated with Texas Instruments. Use at your own risk.
 
-OPTIONS
-
-  --load <file>
-    Load data from <file>
-  --save <file>
-    Save measurement data in <file>
-  --skip <count>
-    Skip <count> data samples. This is useful to avoid startup code
-    influencing the results of a long-running measurement
-  --threshold <watts>|mean
-    Partition data into points with mean power >= <watts> and points with
-    mean power < <watts>, and print some statistics. higher power is handled
-    as peaks, whereas low-power measurements constitute the baseline.
-    If the threshold is set to "mean", the mean power of all measurements
-    will be used
-  --threshold-peakcount <num>
-    Automatically determine a threshold so that there are exactly <num> peaks.
-    A peaks is a group of consecutive measurements with mean power >= threshold.
-    WARNING: In general, there is more than one threshold value leading to
-    exactly <num> peaks. If the difference between baseline and peak
-    power is sufficiently high, this option should do what you mean[tm]
-  --plot=U|I|P
-    Plot voltage / current / power over time
-  --stat
-    Print mean voltage, current, and power as well as total energy consumption.
-  --histogram=<n>
-    Draw histograms of reported energy values per measurement interval
-    (i.e., the differences between each pair of consecutive total energy readings),
-    measurement interval duration, and
-    mean power values per measurement interval
-    (calculated from energy difference and duration).
-    Each histogram uses <n> buckets.
-
 DEPENDENCIES
 
 For data measurements (i.e., any invocation not using --load),
 energytrace-util <https://github.com/carrotIndustries/energytrace-util>
 must be available in $PATH and libmsp430.so must be located in the
 LD library search path (e.g. LD_LIBRARY_PATH=../MSP430Flasher).
-    """
-    )
+
+OPTIONS
+"""
+
+import argparse
+import itertools
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+from shutil import which
+import subprocess
+import sys
+import tempfile
+import time
+
+opt = dict()
 
 
 def running_mean(x: np.ndarray, N: int) -> np.ndarray:
@@ -161,68 +120,85 @@ def peak_search2(data, lower, upper, check_function):
     return None
 
 
-if __name__ == "__main__":
-    try:
-        optspec = "help load= save= skip= threshold= threshold-peakcount= plot= stat histogram="
-        raw_opts, args = getopt.getopt(sys.argv[1:], "", optspec.split(" "))
-
-        for option, parameter in raw_opts:
-            optname = re.sub(r"^--", "", option)
-            opt[optname] = parameter
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter, description=__doc__
+    )
+    parser.add_argument("--load", metavar="FILE", type=str, help="Load data from FILE")
+    parser.add_argument(
+        "--save", metavar="FILE", type=str, help="Save measurement data in FILE"
+    )
+    parser.add_argument(
+        "--skip",
+        metavar="COUNT",
+        type=int,
+        default=0,
+        help="Skip COUNT data samples. This is useful to avoid startup code influencing the results of a long-running measurement",
+    )
+    parser.add_argument(
+        "--threshold",
+        metavar="WATTS",
+        type=str,
+        help="Partition data into points with mean power >= WATTS and points with mean power < WATTS, and print some statistics. higher power is handled as peaks, whereas low-power measurements constitute the baseline. If WATTS is 'mean', the mean power of all measurements will be used",
+    )
+    parser.add_argument(
+        "--threshold-peakcount",
+        metavar="NUM",
+        type=int,
+        help="Automatically determine a threshold so that there are exactly NUM peaks. A peaks is a group of consecutive measurements with mean power >= threshold. WARNING: In general, there is more than one threshold value leading to exactly NUM peaks. If the difference between baseline and peak power is sufficiently high, this option should do what you mean[tm]",
+    )
+    parser.add_argument(
+        "--plot",
+        metavar="UNIT",
+        choices=["U", "I", "P"],
+        help="Plot voltage / current / power over time",
+    )
+    parser.add_argument(
+        "--stat",
+        action="store_true",
+        help="Print mean voltage, current, and power as well as total energy consumption",
+    )
+    parser.add_argument(
+        "--histogram",
+        metavar="N",
+        type=int,
+        help="Draw histograms of reported energy values per measurement interval (i.e., the differences between each pair of consecutive total energy readings), measurement interval duration, and mean power values per measurement interval (calculated from energy difference and duration). Each histogram uses N buckets",
+    )
+    parser.add_argument(
+        "duration", type=int, nargs="?", help="Measurement duration in seconds"
+    )
 
-        if "help" in opt:
-            show_help()
-            sys.exit(0)
+    args = parser.parse_args()
 
-        if not "load" in opt:
-            duration = int(args[0])
+    if args.load is None and args.duration is None:
+        print("Either --load or duration must be specified", file=sys.stderr)
+        sys.exit(1)
 
-        if not "save" in opt:
-            opt["save"] = None
+    if args.threshold is not None and args.threshold != "mean":
+        args.threshold = float(args.threshold)
 
-        if "skip" in opt:
-            opt["skip"] = int(opt["skip"])
-        else:
-            opt["skip"] = 0
-
-        if "threshold" in opt and opt["threshold"] != "mean":
-            opt["threshold"] = float(opt["threshold"])
-
-        if "threshold-peakcount" in opt:
-            opt["threshold-peakcount"] = int(opt["threshold-peakcount"])
-
-    except getopt.GetoptError as err:
-        print(err)
-        sys.exit(2)
-    except IndexError:
-        print("Usage: msp430-etv <duration>")
-        sys.exit(2)
-    except ValueError:
-        print("Error: duration or skip is not a number")
-        sys.exit(2)
-
-    if "load" in opt:
-        if ".xz" in opt["load"]:
+    if args.load:
+        if args.load.endswith(".xz"):
             import lzma
 
-            with lzma.open(opt["load"], "rt") as f:
+            with lzma.open(args.load, "rt") as f:
                 log_data = f.read()
         else:
-            with open(opt["load"], "r") as f:
+            with open(args.load, "r") as f:
                 log_data = f.read()
     else:
-        log_data = measure_data(opt["save"], duration)
+        log_data = measure_data(args.save, args.duration)
 
     lines = log_data.split("\n")
     data_count = sum(map(lambda x: len(x) > 0 and x[0] != "#", lines))
     data_lines = filter(lambda x: len(x) > 0 and x[0] != "#", lines)
 
-    data = np.empty((data_count - opt["skip"], 4))
+    data = np.empty((data_count - args.skip, 4))
 
     energy_overflow_count = 0
     prev_total_energy = 0
     for i, line in enumerate(data_lines):
-        if i >= opt["skip"]:
+        if i >= args.skip:
             fields = line.split(" ")
             if len(fields) == 4:
                 timestamp, current, voltage, total_energy = map(int, fields)
@@ -235,7 +211,7 @@ if __name__ == "__main__":
                 energy_overflow_count += 1
             prev_total_energy = total_energy
             total_energy += energy_overflow_count * (2 ** 32)
-            data[i - opt["skip"]] = [timestamp, current, voltage, total_energy]
+            data[i - args.skip] = [timestamp, current, voltage, total_energy]
 
     m_duration_us = data[-1, 0] - data[0, 0]
     m_energy_nj = data[-1, 3] - data[0, 3]
@@ -255,7 +231,7 @@ if __name__ == "__main__":
         (data[1:, 0] - data[:-1, 0]) * 1e-6
     )
 
-    if "threshold-peakcount" in opt:
+    if args.threshold_peakcount:
         bs_mean = np.mean(power)
 
         # Finding the correct threshold is tricky. If #peaks < peakcont, our
@@ -271,7 +247,7 @@ if __name__ == "__main__":
         # #peaks != peakcount and threshold >= mean, we go down.
         # If that doesn't work, we fall back to a linear search in 1 µW steps
         def direction_function(peakcount, power):
-            if peakcount == opt["threshold-peakcount"]:
+            if peakcount == args.threshold - peakcount:
                 return 0
             if power < bs_mean:
                 return 1
@@ -289,41 +265,41 @@ if __name__ == "__main__":
                     threshold * 1e6, threshold
                 )
             )
-            opt["threshold"] = threshold
+            args.threshold = threshold
         else:
             print("Found no working threshold")
 
-    if "threshold" in opt:
-        if opt["threshold"] == "mean":
-            opt["threshold"] = np.mean(power)
+    if args.threshold:
+        if args.threshold == "mean":
+            args.threshold = np.mean(power)
             print(
                 "Threshold set to {:.0f} µW         : {:.9f}".format(
-                    opt["threshold"] * 1e6, opt["threshold"]
+                    args.threshold * 1e6, args.threshold
                 )
             )
 
         baseline_mean = 0
-        if np.any(power < opt["threshold"]):
-            baseline_mean = np.mean(power[power < opt["threshold"]])
+        if np.any(power < args.threshold):
+            baseline_mean = np.mean(power[power < args.threshold])
             print(
                 "Baseline mean: {:.0f} µW           : {:.9f}".format(
                     baseline_mean * 1e6, baseline_mean
                 )
             )
-        if np.any(power >= opt["threshold"]):
+        if np.any(power >= args.threshold):
             print(
                 "Peak mean: {:.0f} µW               : {:.9f}".format(
-                    np.mean(power[power >= opt["threshold"]]) * 1e6,
-                    np.mean(power[power >= opt["threshold"]]),
+                    np.mean(power[power >= args.threshold]) * 1e6,
+                    np.mean(power[power >= args.threshold]),
                 )
             )
 
         peaks = []
         peak_start = -1
         for i, dp in enumerate(power):
-            if dp >= opt["threshold"] and peak_start == -1:
+            if dp >= args.threshold and peak_start == -1:
                 peak_start = i
-            elif dp < opt["threshold"] and peak_start != -1:
+            elif dp < args.threshold and peak_start != -1:
                 peaks.append((peak_start, i))
                 peak_start = -1
 
@@ -362,7 +338,7 @@ if __name__ == "__main__":
     )
     smooth_power = running_mean(power_from_energy, 10)
 
-    if "stat" in opt:
+    if args.stat:
         mean_voltage = np.mean(data[:, 2] * 1e-3)
         mean_current = np.mean(data[:, 1] * 1e-9)
         mean_power = np.mean(data[:, 1] * data[:, 2] * 1e-12)
@@ -383,8 +359,8 @@ if __name__ == "__main__":
             )
         )
 
-    if "plot" in opt:
-        if opt["plot"] == "U":
+    if args.plot:
+        if args.plot == "U":
             # mV
             (energyhandle,) = plt.plot(
                 data[1:, 0] * 1e-6, data[1:, 2] * 1e-3, "b-", label="U", markersize=1
@@ -398,7 +374,7 @@ if __name__ == "__main__":
             )
             plt.legend(handles=[energyhandle, meanhandle])
             plt.ylabel("Voltage [V]")
-        elif opt["plot"] == "I":
+        elif args.plot == "I":
             print(
                 "Warning: The current reported by energytrace is aggressively smoothed and often inaccurate."
             )
@@ -434,12 +410,12 @@ if __name__ == "__main__":
             plt.ylabel("Power [W]")
         plt.xlabel("Time [s]")
         plt.grid(True)
-        if "load" in opt:
-            plt.title(opt["load"])
+        if args.load:
+            plt.title(args.load)
         plt.show()
 
-    if "histogram" in opt:
-        bin_count = int(opt["histogram"])
+    if args.histogram:
+        bin_count = args.histogram
 
         plt.title("EnergyTrace Data Analysis")
         plt.xlabel("Reported Energy per Measurement Interval [J]")
@@ -464,3 +440,7 @@ if __name__ == "__main__":
         plt.ylabel("Count")
         plt.hist(smooth_power, bins=bin_count)
         plt.show()
+
+
+if __name__ == "__main__":
+    main()