Use StrEnum for units

liskin · Nov 21, 2023 · fc36381 · fc36381
1 parent 67b7041
commit fc36381
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 4 deletions.
diff --git a/src/strava_gear/cli.py b/src/strava_gear/cli.py
@@ -11,6 +11,7 @@
 from .input.activities import read_input_csv
 from .input.activities import read_strava_offline
 from .input.rules import read_rules
+from .report import Units
 from .report import reports
 
 
@@ -54,7 +55,7 @@
     '--show-vert/--hide-vert', default=False, show_default=True,
     help="Show vertical (elevation gain)")
 @click.option(
-    '--units', type=click.Choice(["metric", "imperial"]), default="metric", show_default=True,
+    '--units', type=click.Choice(list(Units)), default=Units.METRIC, show_default=True,
     help="Show data in metric or imperial")
 def main(
     rules_input: TextIO,
@@ -66,7 +67,7 @@ def main(
     show_name: bool,
     show_first_last: bool,
     show_vert: bool,
-    units: str,
+    units: Units,
 ):
     if csv:
         aliases, activities = read_input_csv(csv)

diff --git a/src/strava_gear/report.py b/src/strava_gear/report.py
@@ -1,5 +1,7 @@
 from collections import defaultdict
 import csv
+from enum import StrEnum
+from enum import auto
 from functools import partial
 from typing import Dict
 from typing import Final
@@ -16,13 +18,27 @@
 FOOT_IN_METERS: Final[float] = 0.3048
 
 
-def report(f, res: Result, output, tablefmt: str, show_name: bool, show_first_last: bool, show_vert: bool, units: str):
+class Units(StrEnum):
+    METRIC = auto()
+    IMPERIAL = auto()
+
+
+def report(
+    f,
+    res: Result,
+    output,
+    tablefmt: str,
+    show_name: bool,
+    show_first_last: bool,
+    show_vert: bool,
+    units: Units,
+):
     def cols(d: Dict) -> Dict:
         if not show_name:
             del d["name"]
         if not show_first_last:
             del d["first … last"]
-        if units == "imperial":
+        if units == Units.IMPERIAL:
             del d["km"]
             del d["vert m"]
             if not show_vert: