Refactor fetch_mlr_lineage_fitness script

Small refactor of fetch_mlr_lineage_fitness script primarily to break out a `main()` function and to call this function from `if __name__ == '__main__':`

Author: trvrb

scripts/fetch_mlr_lineage_fitness.py

@@ -1,8 +1,9 @@
-importrequests
+importargparse
 importjson
 importpandasaspd
-importargparse
+importrequests
 importmath
+importsys
 fromaugur.ioimportread_metadata
 fromaugur.utilsimportwrite_json
 
@@ -15,16 +16,6 @@
 # broken for the all-time builds. It would be possible to swap this to key on
 # clade instead, but I think the greater detail of lineages is better in this case
 
-# Set up argument parser
-parser=argparse.ArgumentParser(description="Fetch MLR lineage fitness and match to strain-level metadata")
-parser.add_argument("--metadata", required=True, help="Path to the metadata TSV")
-parser.add_argument("--metadata-id-columns", default=["strain", "name", "Virus name"], nargs="+", help="List of columns to use as identifiers in the metadata file")
-parser.add_argument("--metadata-clade-attribute", default="Nextclade_pango", help="Matched attribute to MLR variants")
-parser.add_argument("--mlr-url", default="https://data.nextstrain.org/files/workflows/forecasts-ncov/gisaid/pango_lineages/global/mlr/latest_results.json", help="URL to fetch the forecasts JSON data.")
-parser.add_argument("--output-node-data", required=True, help="Path to save the output JSON node data.")
-
-args=parser.parse_args()
-
 deffetch_growth_advantages(mlr_url):
 try:
 response=requests.get(mlr_url)
@@ -42,18 +33,30 @@ def fetch_growth_advantages(mlr_url):
 print(f"Error fetching the JSON file: {e}")
 returnNone
 
-try:
-# Fetch the growth advantages
+defmain():
+# Set up argument parser
+parser=argparse.ArgumentParser(description="Fetch MLR lineage fitness and match to strain-level metadata")
+parser.add_argument("--metadata", required=True, help="Path to the metadata TSV")
+parser.add_argument("--metadata-id-columns", default=["strain", "name", "Virus name"], nargs="+", help="List of columns to use as identifiers in the metadata file")
+parser.add_argument("--metadata-clade-attribute", default="Nextclade_pango", help="Matched attribute to MLR variants")
+parser.add_argument("--mlr-url", default="https://data.nextstrain.org/files/workflows/forecasts-ncov/gisaid/pango_lineages/global/mlr/latest_results.json", help="URL to fetch the forecasts JSON data.")
+parser.add_argument("--output-node-data", required=True, help="Path to save the output JSON node data.")
+
+args=parser.parse_args()
+
+# Fetch the remote growth advantages
 growth_advantages=fetch_growth_advantages(args.mlr_url)
 
-# Load the local metadata file
-metadata_file=args.metadata
-metadata=read_metadata(
-metadata_file,
-id_columns=args.metadata_id_columns
- )
+# Load the local metadata
+try:
+metadata=read_metadata(
+args.metadata,
+id_columns=args.metadata_id_columns
+ )
+exceptFileNotFoundErrorase:
+print(f"Error reading metadata file: {e}", file=sys.stderr)
 
-# Match Nextclade_pango entries to the growth advantage
+# Match Nextclade_pango entries in metadata to the fetched growth advantages
 ifgrowth_advantages:
 metadata[args.metadata_clade_attribute] =metadata[args.metadata_clade_attribute].map(growth_advantages)
 else:
@@ -69,7 +72,8 @@ def fetch_growth_advantages(mlr_url):
 # Save node data
 write_json({"nodes": node_data}, args.output_node_data)
 
-exceptFileNotFoundErrorase:
-print(f"Error reading metadata file: {e}", file=sys.stderr)
-exceptExceptionase:
-print(f"An unexpected error occurred: {e}", file=sys.stderr)
+if__name__=='__main__':
+try:
+main()
+exceptExceptionase:
+print(f"An unexpected error occurred: {e}", file=sys.stderr)