Add type hints to pm25

aiddata · Feb 10, 2023 · 9dcb5e4 · 9dcb5e4
1 parent d21ade6
commit 9dcb5e4
Show file tree

Hide file tree

Showing 8 changed files with 23 additions and 77 deletions.
diff --git a/cru_ts/create_env.sh b/cru_ts/create_env.sh
diff --git a/cru_ts/environment.yml b/cru_ts/environment.yml
diff --git a/cru_ts/export_env.sh b/cru_ts/export_env.sh
diff --git a/cru_ts/flow.py b/cru_ts/flow.py
@@ -89,7 +89,9 @@ def cru_ts(
 
     class_instance = CRU_TS(cru_vers, years, raw_dir, output_dir, overwrite_download, overwrite_unzip, overwrite_processing)
 
+
     if task_runner != 'hpc':
         os.chdir(tmp_dir)
-
-    class_instance.run(backend=backend, task_runner=task_runner, run_parallel=run_parallel, max_workers=max_workers, log_dir=timestamp_log_dir, cluster_kwargs=cluster_kwargs)
+        class_instance.run(backend=backend, task_runner=task_runner, run_parallel=run_parallel, max_workers=max_workers, log_dir=timestamp_log_dir)
+    else:
+        class_instance.run(backend=backend, task_runner=task_runner, run_parallel=run_parallel, max_workers=max_workers, log_dir=timestamp_log_dir, cluster=cluster, cluster_kwargs=hpc_cluster_kwargs)
diff --git a/cru_ts/jobscript b/cru_ts/jobscript
diff --git a/cru_ts/main.py b/cru_ts/main.py
@@ -35,14 +35,14 @@
 from configparser import ConfigParser
 from typing import List
 
-
 import rasterio
 import numpy as np
 
 sys.path.insert(1, os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))), 'global_scripts'))
 
 from dataset import Dataset
 
+
 class CRU_TS(Dataset):
 
     name = "Climatic Research Unit gridded Time Series"

diff --git a/pm25/flow.py b/pm25/flow.py
@@ -1,8 +1,8 @@
 import os
-import sys
 from pathlib import Path
 from datetime import datetime
 from configparser import ConfigParser
+from typing import List, Literal
 
 from prefect import flow
 from prefect.filesystems import GitHub
@@ -21,7 +21,20 @@
 
 
 @flow
-def pm25(raw_dir, output_dir, box_config_path, version, years, overwrite_downloads, verify_existing_downloads, overwrite_processing, backend, task_runner, run_parallel, max_workers, log_dir):
+def pm25(
+        raw_dir: str, 
+        output_dir: str, 
+        box_config_path: str, 
+        version: str, 
+        years: List[int], 
+        overwrite_downloads: bool, 
+        verify_existing_downloads: bool, 
+        overwrite_processing: bool, 
+        backend: Literal["local", "mpi", "prefect"],
+        task_runner: Literal["sequential", "concurrent", "dask", "hpc"],
+        run_parallel: bool, 
+        max_workers: int, 
+        log_dir: str):
 
     timestamp = datetime.today()
     time_str = timestamp.strftime("%Y_%m_%d_%H_%M")

diff --git a/pm25/main.py b/pm25/main.py
@@ -88,9 +88,9 @@ def __init__(self,
                  box_config_path: str,
                  version: str,
                  years: list,
-                 overwrite_downloads=True,
-                 verify_existing_downloads=True,
-                 overwrite_processing=False,):
+                 overwrite_downloads: bool,
+                 verify_existing_downloads: bool,
+                 overwrite_processing: bool,):
 
         self.version = version