Appease the type checking gods.

opendatacube · Jul 31, 2024 · ca816a3 · ca816a3
1 parent b6957f6
commit ca816a3
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 7 deletions.
diff --git a/datacube_ows/config_utils.py b/datacube_ows/config_utils.py
@@ -34,7 +34,7 @@
 
 CFG_DICT = dict[str, RAW_CFG]
 
-F = TypeVar('F', bound=Callable[..., Any])
+F = Callable[..., Any]
 
 
 # inclusions defaulting to an empty list is dangerous, but note that it is never modified.

diff --git a/datacube_ows/feature_info.py b/datacube_ows/feature_info.py
@@ -264,7 +264,7 @@ def feature_info(args: dict[str, str]) -> FlaskResponse:
                         # Function signature: pass in:
                         # * a multi-date single pixel (1x1xn) multiband xarray Dataset,
                         date_info[k] = f(data)
-                feature_json["data"].append(date_info)
+                cast(list[RAW_CFG], feature_json["data"]).append(date_info)
         feature_json["data_available_for_dates"] = []
         pt_native = None
         for d in all_time_datasets.coords["time"].values:

diff --git a/datacube_ows/ows_configuration.py b/datacube_ows/ows_configuration.py
@@ -34,7 +34,7 @@
 from slugify import slugify
 
 from datacube_ows.config_utils import (CFG_DICT, RAW_CFG, ConfigException,
-                                       FlagProductBands, FunctionWrapper,
+                                       F, FlagProductBands, FunctionWrapper,
                                        ODCInitException, OWSConfigEntry,
                                        OWSEntryNotFound, OWSExtensibleConfigEntry,
                                        OWSFlagBand, OWSMetadataConfig,
@@ -710,7 +710,7 @@ def parse_feature_info(self, cfg: CFG_DICT):
             _LOG.warning("In layer %s: The 'include_custom' directive is deprecated and will be removed in version 1.9. "
                          "Please refer to the documentation for information on how to migrate your configuration "
                          "to the new 'custom_includes' directive.", self.name)
-        custom = cfg.get("custom_includes", {})
+        custom = cast(dict[str, CFG_DICT | str | F], cfg.get("custom_includes", {}))
         self.feature_info_custom_includes = {k: FunctionWrapper(self, v) for k, v in custom.items()}
 
     # pylint: disable=attribute-defined-outside-init

diff --git a/datacube_ows/styles/base.py b/datacube_ows/styles/base.py
@@ -17,7 +17,7 @@
 
 import datacube_ows.band_utils
 from datacube_ows.config_utils import (CFG_DICT, RAW_CFG, AbstractMaskRule,
-                                       ConfigException, FlagBand,
+                                       ConfigException, FlagBand, F,
                                        FlagProductBands, FunctionWrapper,
                                        OWSConfigEntry, OWSEntryNotFound,
                                        OWSExtensibleConfigEntry,
@@ -204,9 +204,10 @@ def __init__(self, product: "datacube_ows.ows_configuration.OWSNamedLayer",
         self.declare_unready("needed_bands")
         self.declare_unready("flag_bands")
 
+        custom_includes = cast(dict[str, CFG_DICT | str | F], style_cfg.get("custom_includes", {}))
         self.feature_info_includes = {
             k: FunctionWrapper(self, v)
-            for k, v in style_cfg.get("custom_includes", {}).items()
+            for k, v in custom_includes.items()
         }
         self.legend_cfg = self.Legend(self, cast(CFG_DICT, raw_cfg.get("legend", {})))
         if not defer_multi_date:
@@ -539,9 +540,10 @@ def __init__(self, style: "StyleDefBase", cfg: CFG_DICT) -> None:
                     raise ConfigException("Aggregator function is required for non-animated multi-date handlers.")
             self.legend_cfg = self.Legend(self, cast(CFG_DICT, raw_cfg.get("legend", {})))
             self.preserve_user_date_order = cast(bool, cfg.get("preserve_user_date_order", False))
+            custom_includes = cast(dict[str, CFG_DICT | str | F], cfg.get("custom_includes", {}))
             self.feature_info_includes = {
                 k: FunctionWrapper(self.style, v)
-                for k, v in cfg.get("custom_includes", {}).items()
+                for k, v in custom_includes.items()
             }
 
         def applies_to(self, count: int) -> bool:

diff --git a/datacube_ows/styles/ramp.py b/datacube_ows/styles/ramp.py
@@ -589,11 +589,13 @@ def transform_data(self, data: Dataset) -> Dataset:
             :return: RGBA image xarray.  May have a time dimension
             """
             if self.pass_raw_data:
+                assert self.aggregator is not None  # For type-checker
                 agg = self.aggregator(data)
             else:
                 xformed_data = cast("ColorRampDef", self.style).apply_index(data)
                 agg = cast(FunctionWrapper, self.aggregator)(xformed_data)
             return self.color_ramp.apply(agg)
+
         class Legend(RampLegendBase):
             pass