diff --git a/raster_loader/cli/bigquery.py b/raster_loader/cli/bigquery.py index f504fa2..8ab20f7 100644 --- a/raster_loader/cli/bigquery.py +++ b/raster_loader/cli/bigquery.py @@ -4,8 +4,9 @@ import click from functools import wraps, partial -from raster_loader.utils import get_default_table_name +from raster_loader.lib.utils import get_default_table_name from raster_loader.io.bigquery import BigQueryConnection, AccessTokenCredentials +from raster_loader.lib.valuelabels import validate_band_valuelabels def catch_exception(func=None, *, handle=Exception): @@ -101,6 +102,18 @@ def bigquery(args=None): type=int, default=6, ) +@click.option( + "--band-valuelabels", + help="Custom data for valuelabels in JSON format, or 'None'. " + "i.e: '{: , : , ...}'. " + "Could repeat --band-valuelabels to specify multiple bands data. " + "They will be considered in the order they appear in the file. " + "Note that you can set any value to 'None' to omit valuelabels for that band. ", + type=str, + default=[], + multiple=True, + callback=validate_band_valuelabels, +) @catch_exception() def upload( file_path, @@ -119,6 +132,7 @@ def upload( exact_stats=False, basic_stats=False, compression_level=6, + band_valuelabels=[], ): from raster_loader.io.common import ( get_number_of_blocks, @@ -194,6 +208,7 @@ def upload( basic_stats=basic_stats, compress=compress, compression_level=compression_level, + band_valuelabels=band_valuelabels, ) click.echo("Raster file uploaded to Google BigQuery") diff --git a/raster_loader/cli/databricks.py b/raster_loader/cli/databricks.py index c5767de..95c64de 100644 --- a/raster_loader/cli/databricks.py +++ b/raster_loader/cli/databricks.py @@ -4,8 +4,9 @@ import click from functools import wraps, partial -from raster_loader.utils import get_default_table_name +from raster_loader.lib.utils import get_default_table_name from raster_loader.io.databricks import DatabricksConnection +from raster_loader.lib.valuelabels import validate_band_valuelabels def catch_exception(func=None, *, handle=Exception): @@ -115,6 +116,18 @@ def databricks(args=None): type=int, default=6, ) +@click.option( + "--band-valuelabels", + help="Custom data for valuelabels in JSON format, or 'None'. " + "i.e: '{: , : , ...}'. " + "Could repeat --band-valuelabels to specify multiple bands data. " + "They will be considered in the order they appear in the file. " + "Note that you can set any value to 'None' to omit valuelabels for that band. ", + type=str, + default=[], + multiple=True, + callback=validate_band_valuelabels, +) @catch_exception() def upload( server_hostname, @@ -136,6 +149,7 @@ def upload( exact_stats=False, basic_stats=False, compression_level=6, + band_valuelabels=[], ): from raster_loader.io.common import ( get_number_of_blocks, @@ -213,6 +227,7 @@ def upload( basic_stats=basic_stats, compress=compress, compression_level=compression_level, + band_valuelabels=band_valuelabels, ) click.echo("Raster file uploaded to Databricks") diff --git a/raster_loader/cli/snowflake.py b/raster_loader/cli/snowflake.py index d0056eb..d91af69 100644 --- a/raster_loader/cli/snowflake.py +++ b/raster_loader/cli/snowflake.py @@ -4,8 +4,9 @@ import click from functools import wraps, partial -from raster_loader.utils import get_default_table_name, check_private_key +from raster_loader.lib.utils import get_default_table_name, check_private_key from raster_loader.io.snowflake import SnowflakeConnection +from raster_loader.lib.valuelabels import validate_band_valuelabels def catch_exception(func=None, *, handle=Exception): @@ -123,6 +124,18 @@ def snowflake(args=None): type=int, default=6, ) +@click.option( + "--band-valuelabels", + help="Custom data for valuelabels in JSON format, or 'None'. " + "i.e: '{: , : , ...}'. " + "Could repeat --band-valuelabels to specify multiple bands data. " + "They will be considered in the order they appear in the file. " + "Note that you can set any value to 'None' to omit valuelabels for that band. ", + type=str, + default=[], + multiple=True, + callback=validate_band_valuelabels, +) @catch_exception() def upload( account, @@ -148,6 +161,7 @@ def upload( exact_stats=False, basic_stats=False, compression_level=6, + band_valuelabels=[], ): from raster_loader.io.common import ( get_number_of_blocks, @@ -255,6 +269,7 @@ def upload( basic_stats=basic_stats, compress=compress, compression_level=compression_level, + band_valuelabels=band_valuelabels, ) click.echo("Raster file uploaded to Snowflake") diff --git a/raster_loader/io/bigquery.py b/raster_loader/io/bigquery.py index e22951d..d87413d 100644 --- a/raster_loader/io/bigquery.py +++ b/raster_loader/io/bigquery.py @@ -6,8 +6,8 @@ from itertools import chain from raster_loader import __version__ -from raster_loader.errors import import_error_bigquery, IncompatibleRasterException -from raster_loader.utils import ask_yes_no_question, batched +from raster_loader.lib.errors import import_error_bigquery, IncompatibleRasterException +from raster_loader.lib.utils import ask_yes_no_question, batched from raster_loader.io.common import ( check_metadata_is_compatible, get_number_of_blocks, @@ -18,7 +18,7 @@ update_metadata, ) -from typing import Iterable, List, Tuple +from typing import Dict, Iterable, List, Tuple from functools import partial try: @@ -111,6 +111,7 @@ def upload_raster( basic_stats: bool = False, compress: bool = False, compression_level: int = 6, + band_valuelabels: List[Dict[int, str]] = [], ): """Write a raster file to a BigQuery table.""" print("Loading raster file to BigQuery...") @@ -139,6 +140,7 @@ def upload_raster( exact_stats, basic_stats, compress=compress, + band_valuelabels=band_valuelabels, ) overviews_records_gen = rasterio_overview_to_records( diff --git a/raster_loader/io/common.py b/raster_loader/io/common.py index b9e2bb1..fb92037 100644 --- a/raster_loader/io/common.py +++ b/raster_loader/io/common.py @@ -10,16 +10,18 @@ from typing import Dict, Callable, Iterable, List, Tuple, Union from affine import Affine from shapely import wkt # Can not use directly from shapely.wkt +from raster_loader.lib.valuelabels import get_band_valuelabels import rio_cogeo import rasterio import quadbin -from raster_loader.geo import raster_bounds -from raster_loader.errors import ( +from raster_loader.lib.geo import raster_bounds +from raster_loader.lib.errors import ( error_not_google_compatible, ) -from raster_loader.utils import warnings +from raster_loader.lib.utils import warnings + DEFAULT_COG_BLOCK_SIZE = 256 @@ -210,6 +212,7 @@ def rasterio_metadata( exact_stats: bool = False, basic_stats: bool = False, compress: bool = False, + band_valuelabels: List[Dict[int, str]] = [], ): """Open a raster file with rasterio.""" raster_info = rio_cogeo.cog_info(file_path).dict() @@ -278,12 +281,14 @@ def rasterio_metadata( band_nodata = band_value_as_string( raster_dataset, band, band_nodata_value(raster_dataset, band) ) + meta = { "band": band, "type": raster_band_type(raster_dataset, band), "name": band_field_name(band_name, band, band_rename_function), "colorinterp": band_colorinterp, "colortable": get_color_table(raster_dataset, band), + "valuelabels": get_band_valuelabels(file_path, band, band_valuelabels), "stats": stats, "nodata": band_nodata, } @@ -311,6 +316,7 @@ def rasterio_metadata( "colorinterp": e["colorinterp"], "nodata": e["nodata"], "colortable": e["colortable"], + "valuelabels": e["valuelabels"], } for e in bands_metadata ] @@ -506,7 +512,7 @@ def compute_quantiles(data: List[Union[int, float]], cast_function: Callable) -> def get_stats( raster_dataset: rasterio.io.DatasetReader, band: int -) -> rasterio.Statistics: +) -> rasterio._io.Statistics: """Get statistics for a raster band.""" try: # stats method is supported since rasterio 1.4.0 and statistics @@ -937,7 +943,8 @@ def is_valid_raster_dataset(raster_dataset: rasterio.io.DatasetReader) -> bool: def band_without_stats(band): return { k: band[k] - for k in set(list(band.keys())) - set(["stats", "colorinterp", "colortable"]) + for k in set(list(band.keys())) + - set(["stats", "colorinterp", "colortable", "valuelabels"]) } diff --git a/raster_loader/io/databricks.py b/raster_loader/io/databricks.py index 7164dea..a2772fd 100644 --- a/raster_loader/io/databricks.py +++ b/raster_loader/io/databricks.py @@ -3,8 +3,11 @@ import rasterio from itertools import chain -from raster_loader.errors import import_error_databricks, IncompatibleRasterException -from raster_loader.utils import ask_yes_no_question, batched +from raster_loader.lib.errors import ( + import_error_databricks, + IncompatibleRasterException, +) +from raster_loader.lib.utils import ask_yes_no_question, batched from raster_loader.io.common import ( check_metadata_is_compatible, get_number_of_blocks, @@ -15,7 +18,7 @@ update_metadata, ) -from typing import Iterable, List, Tuple +from typing import Dict, Iterable, List, Tuple try: from databricks.connect import DatabricksSession @@ -141,6 +144,7 @@ def upload_raster( basic_stats: bool = False, compress: bool = False, compression_level: int = 6, + band_valuelabels: List[Dict[int, str]] = [], ): """Write a raster file to a Databricks table.""" # Wait for cluster to be ready before starting the upload @@ -174,6 +178,7 @@ def upload_raster( exact_stats, basic_stats, compress=compress, + band_valuelabels=band_valuelabels, ) overviews_records_gen = rasterio_overview_to_records( diff --git a/raster_loader/io/snowflake.py b/raster_loader/io/snowflake.py index 6ecd9c1..04699b6 100644 --- a/raster_loader/io/snowflake.py +++ b/raster_loader/io/snowflake.py @@ -3,14 +3,14 @@ import pandas as pd from itertools import chain -from typing import Iterable, List, Tuple +from typing import Dict, Iterable, List, Tuple -from raster_loader.errors import ( +from raster_loader.lib.errors import ( IncompatibleRasterException, import_error_snowflake, ) -from raster_loader.utils import ask_yes_no_question, batched +from raster_loader.lib.utils import ask_yes_no_question, batched from raster_loader.io.common import ( rasterio_metadata, @@ -208,6 +208,7 @@ def upload_raster( basic_stats: bool = False, compress: bool = False, compression_level: int = 6, + band_valuelabels: List[Dict[int, str]] = [], ) -> bool: """Write a raster file to a Snowflake table.""" @@ -243,6 +244,7 @@ def band_rename_function(x): exact_stats, basic_stats, compress, + band_valuelabels, ) overviews_records_gen = rasterio_overview_to_records( diff --git a/raster_loader/lib/__init__.py b/raster_loader/lib/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/raster_loader/errors.py b/raster_loader/lib/errors.py similarity index 100% rename from raster_loader/errors.py rename to raster_loader/lib/errors.py diff --git a/raster_loader/geo.py b/raster_loader/lib/geo.py similarity index 100% rename from raster_loader/geo.py rename to raster_loader/lib/geo.py diff --git a/raster_loader/utils.py b/raster_loader/lib/utils.py similarity index 100% rename from raster_loader/utils.py rename to raster_loader/lib/utils.py diff --git a/raster_loader/lib/valuelabels.py b/raster_loader/lib/valuelabels.py new file mode 100644 index 0000000..3b5d56b --- /dev/null +++ b/raster_loader/lib/valuelabels.py @@ -0,0 +1,25 @@ +import json +import click +from typing import Dict, List, Optional + + +def get_band_valuelabels( + file_path: str, band: int, band_valuelabels: List[Dict[int, str]] +) -> Optional[Dict[int, str]]: + valuelabels = None + if len(band_valuelabels) >= band and band_valuelabels[band - 1] is not None: + print(f"Using the provided valuelabels for band {band}") + valuelabels = band_valuelabels[band - 1] + return valuelabels + + +def validate_band_valuelabels(_, __, value): + """ + Validate the band valuelabels parameter for click library callback. + """ + try: + return [json.loads(item) if item != "None" else None for item in value] + except json.JSONDecodeError: + raise click.BadParameter( + "Invalid JSON format. Please provide a valid JSON object." + ) diff --git a/raster_loader/tests/bigquery/test_cli.py b/raster_loader/tests/bigquery/test_cli.py index 5478e94..bb5d5b8 100644 --- a/raster_loader/tests/bigquery/test_cli.py +++ b/raster_loader/tests/bigquery/test_cli.py @@ -308,6 +308,33 @@ def test_bigquery_describe(*args, **kwargs): assert result.exit_code == 0 +@patch( + "raster_loader.cli.bigquery.BigQueryConnection.upload_records", return_value=None +) +@patch("raster_loader.cli.bigquery.BigQueryConnection.__init__", return_value=None) +def test_bigquery_upload_custom_valuelabels(*args, **kwargs): + runner = CliRunner() + result = runner.invoke( + main, + [ + "bigquery", + "upload", + "--file_path", + f"{tiff}", + "--project", + "project", + "--dataset", + "dataset", + "--table", + "table", + "--band-valuelabels", + '{"90": "Woody Wetlands", "52": "Shrub/Scrub"}', + ], + ) + assert result.exit_code == 1 + assert "Using the provided valuelabels for band 1" in result.output + + def test_info(*args, **kwargs): runner = CliRunner() result = runner.invoke(main, ["info"]) diff --git a/raster_loader/tests/bigquery/test_io.py b/raster_loader/tests/bigquery/test_io.py index f7a9e45..b982b76 100644 --- a/raster_loader/tests/bigquery/test_io.py +++ b/raster_loader/tests/bigquery/test_io.py @@ -483,6 +483,7 @@ def test_rasterio_to_table_overwrite(*args, **kwargs): }, "nodata": "0", "colortable": None, + "valuelabels": None, } ], "num_blocks": 1, @@ -649,6 +650,7 @@ def test_rasterio_to_table_invalid_raster(*args, **kwargs): }, "nodata": "0", "colortable": None, + "valuelabels": None, } ], "num_blocks": 1, @@ -744,6 +746,7 @@ def test_get_labels(*args, **kwargs): }, "nodata": "0", "colortable": None, + "valuelabels": None, } ], "num_blocks": 1, @@ -802,6 +805,7 @@ def test_rasterio_to_bigquery_with_compression(*args, **kwargs): }, "nodata": "0", "colortable": None, + "valuelabels": None, } ], "num_blocks": 1, diff --git a/raster_loader/tests/databricks/test_io.py b/raster_loader/tests/databricks/test_io.py index 1991bdb..82c34f7 100644 --- a/raster_loader/tests/databricks/test_io.py +++ b/raster_loader/tests/databricks/test_io.py @@ -391,6 +391,7 @@ def test_rasterio_to_table(*args, **kwargs): "nodata": "0", "colorinterp": "red", "colortable": None, + "valuelabels": None, } ], "num_blocks": 1, diff --git a/raster_loader/tests/snowflake/test_io.py b/raster_loader/tests/snowflake/test_io.py index d38c9e4..0437b7d 100644 --- a/raster_loader/tests/snowflake/test_io.py +++ b/raster_loader/tests/snowflake/test_io.py @@ -462,6 +462,7 @@ def test_rasterio_to_table_overwrite(*args, **kwargs): "nodata": "0", "colorinterp": "red", "colortable": None, + "valuelabels": None, } ], "num_blocks": 1, @@ -629,6 +630,7 @@ def test_rasterio_to_table_invalid_raster(*args, **kwargs): "nodata": "0", "colorinterp": "red", "colortable": None, + "valuelabels": None, } ], "num_blocks": 1, @@ -712,6 +714,7 @@ def test_append_with_different_resolution(*args, **kwargs): "nodata": "0", "colorinterp": "red", "colortable": None, + "valuelabels": None, } ], "num_blocks": 1, @@ -772,6 +775,7 @@ def test_rasterio_to_snowflake_with_compression(*args, **kwargs): "nodata": "0", "colorinterp": "red", "colortable": None, + "valuelabels": None, } ], "num_blocks": 1, diff --git a/raster_loader/tests/test_utils.py b/raster_loader/tests/test_utils.py index 24dfdde..5b1e892 100644 --- a/raster_loader/tests/test_utils.py +++ b/raster_loader/tests/test_utils.py @@ -1,7 +1,7 @@ import builtins from unittest.mock import patch -from raster_loader.utils import ask_yes_no_question +from raster_loader.lib.utils import ask_yes_no_question def test_ask_yes_no_question_answer_yes(): diff --git a/setup.cfg b/setup.cfg index f0c293b..33bdf7b 100644 --- a/setup.cfg +++ b/setup.cfg @@ -69,6 +69,6 @@ all = %(databricks)s [flake8] -max-line-length = 88 +max-line-length = 100 ignore = E203 W503