obr: add subcommands list and report

globaldothealth · Aug 29, 2024 · c7c8afc · c7c8afc
1 parent b2b32f1
commit c7c8afc
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 20 deletions.
diff --git a/src/obr/__init__.py b/src/obr/__init__.py
@@ -15,30 +15,45 @@ def main():
     parser = argparse.ArgumentParser(
         description="Global.health outbreak report creator"
     )
-    parser.add_argument("outbreak", help="Outbreak name")
-    parser.add_argument("url", help="Data URL")
-    parser.add_argument("-b", "--bucket", help="S3 bucket to write outbreak report to")
-    parser.add_argument(
+
+    subparsers = parser.add_subparsers(dest="command")
+
+    report_parser = subparsers.add_parser("report", help="Generate briefing report")
+    _ = subparsers.add_parser("list", help="List outbreaks known to obr")
+    report_parser.add_argument("outbreak", help="Outbreak name")
+    report_parser.add_argument("url", help="Data URL")
+    report_parser.add_argument(
+        "-b", "--bucket", help="S3 bucket to write outbreak report to"
+    )
+    report_parser.add_argument(
         "--cloudfront", help="Cloudfront distribution which should be invalidated"
     )
-    parser.add_argument(
+    report_parser.add_argument(
         "-o", "--open", action="store_true", help="Open local file in web browser"
     )
 
     args = parser.parse_args()
-    if args.outbreak not in OUTBREAKS:
-        abort(f"Outbreak not supported: {args.outbreak}")
-
-    build(
-        args.outbreak,
-        args.url,
-        OUTBREAKS[args.outbreak]["plots"],
-        date_columns=OUTBREAKS[args.outbreak].get("additional_date_columns", []),
-        output_bucket=args.bucket,
-        cloudfront_distribution=args.cloudfront,
-    )
-    if args.open and (Path(args.outbreak + ".html")).exists():
-        webbrowser.open("file://" + str(Path.cwd() / (args.outbreak + ".html")))
+    match args.command:
+        case "list":
+            for outbreak in OUTBREAKS:
+                print(
+                    f"\033[1m{outbreak:12s} \033[0m{OUTBREAKS[outbreak]["description"]}"
+                )
+        case "report":
+            if args.outbreak not in OUTBREAKS:
+                abort(f"Outbreak not supported: {args.outbreak}")
+            build(
+                args.outbreak,
+                args.url,
+                OUTBREAKS[args.outbreak]["plots"],
+                date_columns=OUTBREAKS[args.outbreak].get(
+                    "additional_date_columns", []
+                ),
+                output_bucket=args.bucket,
+                cloudfront_distribution=args.cloudfront,
+            )
+            if args.open and (Path(args.outbreak + ".html")).exists():
+                webbrowser.open("file://" + str(Path.cwd() / (args.outbreak + ".html")))
 
 
 if __name__ == "__main__":

diff --git a/src/obr/outbreaks/__init__.py b/src/obr/outbreaks/__init__.py
@@ -87,7 +87,14 @@
 ]
 
 OUTBREAKS = {
-    "marburg": {"plots": outbreak_marburg, "additional_date_columns": ["Data_up_to"]},
-    "mpox-2024": {"plots": outbreak_mpox_2024},
+    "marburg": {
+        "description": "Marburg 2023 [GHL2023.D11.1D60.1]",
+        "plots": outbreak_marburg,
+        "additional_date_columns": ["Data_up_to"],
+    },
+    "mpox-2024": {
+        "description": "Mpox 2024 [GHL2024.D11.1E71]",
+        "plots": outbreak_mpox_2024,
+    },
 }
 __all__ = ["OUTBREAKS"]