linting, typing, etc

GNS-Science · Jul 24, 2023 · 90ff727 · 90ff727
1 parent cbe92a3
commit 90ff727
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 153 deletions.
diff --git a/tests/fixtures/sample_task_args.py b/tests/fixtures/sample_task_args.py
diff --git a/toshi_hazard_post/hazard_aggregation/aws_deaggregation.py b/toshi_hazard_post/hazard_aggregation/aws_deaggregation.py
@@ -76,7 +76,7 @@ def tasks_by_chunk(
 
     count = 0
     total = 0
-    task_chunk = {key: [] for key in keys}
+    task_chunk: Dict[str, Any] = {key: [] for key in keys}
     n_combs = len(locations) * len(aggs) * len(poes) * len(imts) * len(vs30s)
 
     for (location, agg, poe, imt, vs30) in itertools.product(locations, aggs, poes, imts, vs30s):

diff --git a/toshi_hazard_post/hazard_aggregation/deaggregation.py b/toshi_hazard_post/hazard_aggregation/deaggregation.py
@@ -7,7 +7,7 @@
 from collections import namedtuple
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, Generator, List, Tuple, Union
+from typing import Any, Dict, Generator, Iterable, List, Tuple, Union
 
 from nzshm_common.location.code_location import CodedLocation
 from nzshm_common.util import decompress_string
@@ -135,6 +135,7 @@ def requested_configs(
     iter_method: str = '',
 ) -> Generator[DeaggConfig, None, None]:
 
+    iterator: Iterable[Any] = []
     if not iter_method or iter_method.lower() == 'product':
         iterator = itertools.product(
             map(coded_location, locations),

diff --git a/toshi_hazard_post/hazard_grid/aws_gridded_hazard.py b/toshi_hazard_post/hazard_grid/aws_gridded_hazard.py
@@ -1,6 +1,6 @@
 import itertools
 from dataclasses import asdict
-from typing import Dict, Generator, Iterable
+from typing import Any, Collection, Dict, Generator, Iterable
 
 import boto3
 from nzshm_common.location import CodedLocation
@@ -44,10 +44,10 @@ def batch_job_config(task_arguments: Dict, job_arguments: Dict, task_id: int):
 
 def tasks_by_chunk(
     poe_levels: Iterable[float],
-    hazard_model_ids: Iterable[str],
-    vs30s: Iterable[float],
-    imts: Iterable[str],
-    aggs: Iterable[str],
+    hazard_model_ids: Collection[str],
+    vs30s: Collection[float],
+    imts: Collection[str],
+    aggs: Collection[str],
     chunk_size: int,
     iter_method: str,
 ) -> Generator[DistributedGridTaskArguments, None, None]:
@@ -65,6 +65,7 @@ def tasks_by_chunk(
         force=False,
     )
 
+    iterator: Iterable[Any] = []
     if iter_method == 'product':
         iterator = itertools.product(hazard_model_ids, vs30s, imts, aggs)
         total_jobs = len(hazard_model_ids) * len(vs30s) * len(imts) * len(aggs)
@@ -101,10 +102,10 @@ def tasks_by_chunk(
 def batch_job_configs(
     location_grid_id: str,
     poe_levels: Iterable[float],
-    hazard_model_ids: Iterable[str],
-    vs30s: Iterable[float],
-    imts: Iterable[str],
-    aggs: Iterable[str],
+    hazard_model_ids: Collection[str],
+    vs30s: Collection[float],
+    imts: Collection[str],
+    aggs: Collection[str],
     force: bool = False,
     filter_locations: Iterable[CodedLocation] = None,
     iter_method: str = 'product',
@@ -141,10 +142,10 @@ def batch_job_configs(
 def distribute_gridded_hazard(
     location_grid_id: str,
     poe_levels: Iterable[float],
-    hazard_model_ids: Iterable[str],
-    vs30s: Iterable[float],
-    imts: Iterable[str],
-    aggs: Iterable[str],
+    hazard_model_ids: Collection[str],
+    vs30s: Collection[float],
+    imts: Collection[str],
+    aggs: Collection[str],
     force: bool = False,
     filter_locations: Iterable[CodedLocation] = None,
     iter_method: str = 'product',

diff --git a/toshi_hazard_post/hazard_grid/gridded_hazard.py b/toshi_hazard_post/hazard_grid/gridded_hazard.py
@@ -5,7 +5,7 @@
 import multiprocessing
 from collections import namedtuple
 from dataclasses import dataclass
-from typing import Dict, Iterable, List
+from typing import Any, Dict, Iterable, List
 
 import numpy as np
 from nzshm_common.grids import RegionGrid
@@ -169,6 +169,7 @@ def calc_gridded_hazard(
     for w in workers:
         w.start()
 
+    iterator: Iterable[Any] = []
     if iter_method == 'product':
         iterator = itertools.product(hazard_model_ids, vs30s, imts, aggs)
     elif iter_method == 'zip':

diff --git a/toshi_hazard_post/logic_tree/branch_combinator.py b/toshi_hazard_post/logic_tree/branch_combinator.py
@@ -2,7 +2,7 @@
 import logging
 import time
 from pathlib import Path
-from typing import Dict, Iterable, List, Union
+from typing import Collection, Dict, List, Union
 
 # from toshi_hazard_store.query_v3 import get_hazard_metadata_v3
 import toshi_hazard_store
@@ -16,7 +16,7 @@
 log = logging.getLogger(__name__)
 
 
-def preload_meta(ids: Iterable[str], vs30: int) -> Dict[str, dict]:
+def preload_meta(ids: Collection[str], vs30: int) -> Dict[str, dict]:
     """Retreive the GMCM logic tree metadata from Toshi-Hazard-Store.
 
     Parameters