Add raster compression option (#106)

* adding compression option * Adding modification and verification in method * Add argument in method * rename and add docstring * Add `raster_compression` type in docstring * fix None bug and rename * change argument * CodeQL issue * pre-commit issues * Address CodeQL alert * test with pre-commit * change in metadata * add pre-commit * pre-commit change * modify `get_raster_compression` docstring * modify line length on docstring * Problem with self and variables env * Adding conftest and test * rename filepath for authfile * fix authfile fixture * try and fix authfile writing in Github action * modify compression methods * fix authfile * fix authfile * add assert for authentification * remove assert * change None check and add test for comp * Add comments * comply with github_adv_sec * fix typo * forgot pre-commit * add Antsalacia in author list --------- Co-authored-by: Adrien Wehrlé <[email protected]>
AdrienWehrle · Dec 17, 2024 · 15d3c93 · 15d3c93
1 parent f8e5528
commit 15d3c93
Show file tree

Hide file tree

Showing 3 changed files with 515 additions and 383 deletions.
diff --git a/earthspy/earthspy.py b/earthspy/earthspy.py
@@ -1,6 +1,6 @@
 """
 
-@author: Adrien Wehrlé, EO-IO, University of Zurich, Switzerland
+@authors: Adrien Wehrlé (EO-IO), Antsalacia
 
 """
 
@@ -84,6 +84,7 @@ def set_query_parameters(
         download_mode: str = "SM",
         remove_splitboxes: bool = True,
         verbose: bool = True,
+        raster_compression: str = None,
     ) -> None:
         """Define a set of parameters used for the API request.
 
@@ -144,6 +145,11 @@ def set_query_parameters(
         :param verbose: Whether to print processing status or not, defaults
           to True.
         :type verbose: bool, optional
+
+
+        :param raster_compression: Raster compression to apply following methods
+          available in rasterio, defaults to None.
+        :type raster_compression: Union[None, str], optional
         """
 
         # set processing attributes
@@ -171,6 +177,9 @@ def set_query_parameters(
         # set and correct resolution
         self.set_correct_resolution()
 
+        # set compression method
+        self.get_raster_compression(raster_compression)
+
         # set post-processing attributes
         self.get_evaluation_script(evaluation_script)
         self.get_store_folder(store_folder)
@@ -187,6 +196,28 @@ def set_query_parameters(
 
         return None
 
+    def get_raster_compression(self, raster_compression: Union[None, str]) -> str:
+        """Get raster compression based on rasterio's available methods
+
+        :return: Raster compression method
+        :rtype: Union[None, str]
+        """
+
+        # list rasterio compression algorithm and exclude dunders
+        rasterio_compression_algorithms = [
+            m for m in dir(rasterio.enums.Compression) if not m.startswith("__")
+        ]
+
+        # use rasterio compression method as is
+        if raster_compression is None:
+            self.raster_compression = None
+        elif raster_compression.lower() in rasterio_compression_algorithms:
+            self.raster_compression = raster_compression
+        else:
+            raise KeyError("Compression algorithm not found")
+
+        return self.raster_compression
+
     def get_data_collection(self) -> shb.DataCollection:
         """Get Sentinel Hub DataCollection object from data collection name.
 
@@ -294,7 +325,6 @@ def get_raw_data_collection_resolution(self) -> int:
 
         :return: Data collection resolution.
         :rtype: int
-
         """
 
         # set default satellite resolution
@@ -1072,11 +1102,19 @@ def merge_rasters(self) -> None:
                         "transform": output_transform,
                     }
                 )
+
+                # update dictionary if compression set
+                if self.raster_compression is not None:
+                    output_meta.update({"compress": self.raster_compression})
+
+                # extract scene id
                 id_dict = {k: self.metadata[date][0][k] for k in ["id"]}
+
                 # write mosaic
                 with rasterio.open(date_output_filename, "w", **output_meta) as dst:
                     dst.write(mosaic)
                     dst.update_tags(**id_dict)
+
                 # save file name of merged raster
                 self.output_filenames_renamed.append(date_output_filename)
 

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+"""
+
+@authors: Adrien Wehrlé (EO-IO), Antsalacia
+
+"""
+
+import os
+
+import pytest
+
+import earthspy.earthspy as es
+
+
+def pytest_addoption(parser):
+    """Add option to pass local path to authentication file"""
+    parser.addoption(
+        "--authfile",
+        action="store",
+        default="./auth.txt",
+        help="Full path to Sentinel Hub credential file containing ID and password",
+    )
+
+
+# if running in Github action
+if os.getenv("CI") is not None:
+
+    @pytest.fixture(scope="session")
+    def SH_CLIENT_ID() -> None:
+        """Create local client id from environment variable"""
+        # check if variable in environment variables
+        SH_CLIENT_ID = os.environ["SH_CLIENT_ID"]
+        return SH_CLIENT_ID
+
+    @pytest.fixture(scope="session")
+    def SH_CLIENT_SECRET() -> None:
+        """Create local client secret from environment variable"""
+        # check if variable in environment variables
+        SH_CLIENT_SECRET = os.environ["SH_CLIENT_SECRET"]
+        return SH_CLIENT_SECRET
+
+    # path to credential file to be created
+    @pytest.fixture(scope="session")
+    def authfile(SH_CLIENT_ID, SH_CLIENT_SECRET):
+        """Set credential file name and create credential file
+        for testing"""
+        authfile = "auth.txt"
+        with open(authfile, "w") as out:
+            out.write(f"{SH_CLIENT_ID}\n{SH_CLIENT_SECRET}")
+        return authfile
+
+
+# if running locally
+else:
+
+    @pytest.fixture(scope="session")
+    def authfile(pytestconfig):
+        """Get option from command line call"""
+        return pytestconfig.getoption("authfile")
+
+    @pytest.fixture(scope="session")
+    def credentials(authfile):
+        """Read credentials stored in text file"""
+
+        with open(authfile) as file:
+            credentials = file.read().splitlines()
+        return credentials
+
+    @pytest.fixture(scope="session")
+    def SH_CLIENT_ID(credentials) -> None:
+        """Extract client id from line"""
+        SH_CLIENT_ID = credentials[0]
+        return SH_CLIENT_ID
+
+    @pytest.fixture(scope="session")
+    def SH_CLIENT_SECRET(credentials) -> None:
+        """Extract client secret from line"""
+        SH_CLIENT_SECRET = credentials[1]
+        return SH_CLIENT_SECRET
+
+
+@pytest.fixture(scope="session")
+def test_evalscript():
+    """Set a test evalscript for Sentinel-2"""
+    test_evalscript = """
+        //VERSION=3
+        function setup(){
+          return{
+            input: ["B02", "B03", "B04", "dataMask"],
+            output: {bands: 4}
+          }
+        }
+        function evaluatePixel(sample){
+          // Set gain for visualisation
+          let gain = 2.5;
+          // Return RGB
+          return [sample.B04 * gain, sample.B03 * gain, sample.B02 * gain,
+                  sample.dataMask];
+        }
+        """
+    return test_evalscript
+
+
+@pytest.fixture(scope="session")
+def test_url():
+    """Set a test evalscript pointing to Sentinel-2 True Color"""
+    test_url = (
+        "https://custom-scripts.sentinel-hub.com/custom-scripts/"
+        + "sentinel-2/true_color/script.js"
+    )
+    return test_url
+
+
+@pytest.fixture(scope="session")
+def test_collection():
+    """Set a test data collection"""
+    test_collection = "SENTINEL2_L2A"
+    return test_collection
+
+
+@pytest.fixture(scope="session")
+def test_bounding_box():
+    """Set a test footprint area (bounding box)"""
+    test_bounding_box = [-51.13, 69.204, -51.06, 69.225]
+    return test_bounding_box
+
+
+@pytest.fixture(scope="session")
+def test_area_name():
+    """Set a test area available as geojson file"""
+    test_area_name = "Ilulissat"
+    return test_area_name
+
+
+@pytest.fixture(scope="session")
+def t1(authfile, test_evalscript, test_collection, test_bounding_box):
+    """Set a test query with default parameters"""
+    t1 = es.EarthSpy(authfile)
+    t1.set_query_parameters(
+        bounding_box=test_bounding_box,
+        time_interval=["2019-08-23"],
+        evaluation_script=test_evalscript,
+        data_collection=test_collection,
+        download_mode="SM",
+    )
+    return t1
+
+
+@pytest.fixture(scope="session")
+def t2(authfile, test_evalscript, test_collection, test_area_name):
+    """Set a test query with area name"""
+    t2 = es.EarthSpy(authfile)
+    t2.set_query_parameters(
+        bounding_box=test_area_name,
+        time_interval=["2019-08-23"],
+        evaluation_script=test_evalscript,
+        data_collection=test_collection,
+        download_mode="SM",
+    )
+    return t2
+
+
+@pytest.fixture(scope="session")
+def t3(authfile, test_evalscript, test_collection, test_bounding_box):
+    """Set a test query with direct download mode"""
+    t3 = es.EarthSpy(authfile)
+    t3.set_query_parameters(
+        bounding_box=test_bounding_box,
+        time_interval=["2019-08-23"],
+        evaluation_script=test_evalscript,
+        data_collection=test_collection,
+        download_mode="D",
+    )
+    return t3
+
+
+@pytest.fixture(scope="session")
+def t4(authfile, test_evalscript, test_collection, test_bounding_box):
+    """Set a test query with LZW raster compression"""
+    t4 = es.EarthSpy(authfile)
+    t4.set_query_parameters(
+        bounding_box=test_bounding_box,
+        time_interval=["2019-08-23"],
+        evaluation_script=test_evalscript,
+        data_collection=test_collection,
+        download_mode="SM",
+        raster_compression="LZW",
+    )
+    return t4