From 2bbd421af4fe62e2c803d4c43fef9d52fdf11765 Mon Sep 17 00:00:00 2001
From: Stefaan Lippens <stefaan.lippens@vito.be>
Date: Mon, 11 Mar 2024 10:01:22 +0100
Subject: [PATCH] to_scl_dilation_mask: define as DriverDataCube method

makes things easier to track

related to https://github.com/Open-EO/openeo-geopyspark-driver/issues/715
---
 openeo_driver/ProcessGraphDeserializer.py | 32 ++++++++++-------------
 openeo_driver/datacube.py                 | 11 ++++++++
 openeo_driver/dry_run.py                  |  8 ++++--
 3 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/openeo_driver/ProcessGraphDeserializer.py b/openeo_driver/ProcessGraphDeserializer.py
index 8ee85aeb..0d5d2940 100644
--- a/openeo_driver/ProcessGraphDeserializer.py
+++ b/openeo_driver/ProcessGraphDeserializer.py
@@ -1995,24 +1995,20 @@ def mask_scl_dilation(args: Dict, env: EvalEnv):
 @process_registry_2xx.add_function(spec=read_spec("openeo-processes/experimental/to_scl_dilation_mask.json"))
 def to_scl_dilation_mask(args: ProcessArgs, env: EvalEnv):
     cube: DriverDataCube = args.get_required("data", expected_type=DriverDataCube)
-
-    if hasattr(cube, "to_scl_dilation_mask"):
-        # Get default values from spec
-        spec = read_spec("openeo-processes/experimental/to_scl_dilation_mask.json")
-        defaults = {param["name"]: param["default"] for param in spec["parameters"] if "default" in param}
-        optionals = {
-            arg: args.get_optional(arg, default=defaults[arg])
-            for arg in [
-                "erosion_kernel_size",
-                "mask1_values",
-                "mask2_values",
-                "kernel1_size",
-                "kernel2_size",
-            ]
-        }
-        return cube.to_scl_dilation_mask(**optionals)
-    else:
-        raise FeatureUnsupportedException(message="to_scl_dilation_mask is not supported")
+    # Get default values for other args from spec
+    spec = read_spec("openeo-processes/experimental/to_scl_dilation_mask.json")
+    defaults = {param["name"]: param["default"] for param in spec["parameters"] if "default" in param}
+    optionals = {
+        arg: args.get_optional(arg, default=defaults[arg])
+        for arg in [
+            "erosion_kernel_size",
+            "mask1_values",
+            "mask2_values",
+            "kernel1_size",
+            "kernel2_size",
+        ]
+    }
+    return cube.to_scl_dilation_mask(**optionals)
 
 
 @process_registry_100.add_function(spec=read_spec("openeo-processes/experimental/mask_l1c.json"))
diff --git a/openeo_driver/datacube.py b/openeo_driver/datacube.py
index deabaca0..aa9f3fc0 100644
--- a/openeo_driver/datacube.py
+++ b/openeo_driver/datacube.py
@@ -222,6 +222,17 @@ def resample_spatial(
     ):
         self._not_implemented()
 
+    def to_scl_dilation_mask(
+        self,
+        erosion_kernel_size: int,
+        mask1_values: List[int],
+        mask2_values: List[int],
+        kernel1_size: int,
+        kernel2_size: int,
+    ) -> DriverDataCube:
+        # Note: this is a non-standard process
+        self._not_implemented()
+
 
 class VectorCubeError(InternalException):
     code = "VectorCubeError"
diff --git a/openeo_driver/dry_run.py b/openeo_driver/dry_run.py
index b5775a5d..53d4dfad 100644
--- a/openeo_driver/dry_run.py
+++ b/openeo_driver/dry_run.py
@@ -32,6 +32,7 @@
 These source constraints can then be fetched from the EvalEnv at `load_collection` time.
 
 """
+from __future__ import annotations
 import logging
 from enum import Enum
 from typing import List, Union, Tuple, Any, Optional
@@ -737,11 +738,14 @@ def atmospheric_correction(
     def mask_scl_dilation(self, **kwargs) -> 'DriverDataCube':
         return self._process("custom_cloud_mask", arguments={**{"method":"mask_scl_dilation"},**kwargs})
 
-    def to_scl_dilation_mask(self, erosion_kernel_size: int,
+    def to_scl_dilation_mask(
+        self,
+        erosion_kernel_size: int,
         mask1_values: List[int],
         mask2_values: List[int],
         kernel1_size: int,
-        kernel2_size: int) -> 'DriverDataCube':
+        kernel2_size: int,
+    ) -> DryRunDataCube:
         cube = self._process("process_type", [ProcessType.FOCAL_SPACE])
         size = kernel2_size
         cube = cube._process("pixel_buffer", arguments={"buffer_size": [size/2.0,size/2.0]})