rio-tiler 6.7.0__py3-none-any.whl → 7.0.0__py3-none-any.whl

rio_tiler/__init__.py

@@ -1,6 +1,6 @@
 """rio-tiler."""
 
-__version__ = "6.7.0"
+__version__ = "7.0.0"
 
 from . import (  # noqa
     colormap,

rio_tiler/colormap.py

@@ -1,5 +1,6 @@
 """rio-tiler colormap functions and classes."""
 
+import json
 import os
 import pathlib
 import re
@@ -22,23 +23,30 @@ from rio_tiler.types import (
 )
 
 try:
-    from importlib.resources import files as resources_files  # type: ignore
+    from importlib.resources import as_file
+    from importlib.resources import files as resources_files
 except ImportError:
     # Try backported to PY<39 `importlib_resources`.
+    from importlib_resources import as_file  # type: ignore
     from importlib_resources import files as resources_files  # type: ignore
 
 
 EMPTY_COLORMAP: GDALColorMapType = {i: (0, 0, 0, 0) for i in range(256)}
 
-DEFAULT_CMAPS_FILES = {
-    f.stem: str(f)
-    for f in (resources_files(__package__) / "cmap_data").glob("*.npy")  # type: ignore
-}
+_RIO_CMAP_DIR = resources_files(__package__) / "cmap_data"
+with as_file(_RIO_CMAP_DIR) as p:
+    DEFAULT_CMAPS_FILES = {
+        f.stem: f for f in p.glob("**/*") if f.suffix in {".npy", ".json"}
+    }
 
 USER_CMAPS_DIR = os.environ.get("COLORMAP_DIRECTORY", None)
 if USER_CMAPS_DIR:
     DEFAULT_CMAPS_FILES.update(
-        {f.stem: str(f) for f in pathlib.Path(USER_CMAPS_DIR).glob("*.npy")}
+        {
+            f.stem: f
+            for f in pathlib.Path(USER_CMAPS_DIR).glob("**/*")
+            if f.suffix in {".npy", ".json"}
+        }
     )
 
 
@@ -129,7 +137,7 @@ def apply_discrete_cmap(data: numpy.ndarray, colormap: GDALColorMapType) -> Data
 
     Args:
         data (numpy.ndarray): 1D image array to translate to RGB.
-        color_map (dict): Discrete ColorMap dictionary.
+        colormap (GDALColorMapType): Discrete ColorMap dictionary.
 
     Returns:
         tuple: Data (numpy.ndarray) and Alpha band (numpy.ndarray).
@@ -168,7 +176,7 @@ def apply_intervals_cmap(
 
     Args:
         data (numpy.ndarray): 1D image array to translate to RGB.
-        color_map (Sequence): Sequence of intervals and color in form of [([min, max], [r, g, b, a]), ...].
+        colormap (IntervalColorMapType): Sequence of intervals and color in form of [([min, max], [r, g, b, a]), ...].
 
     Returns:
         tuple: Data (numpy.ndarray) and Alpha band (numpy.ndarray).
@@ -274,7 +282,7 @@ class ColorMaps:
 
     """
 
-    data: Dict[str, Union[str, ColorMapType]] = attr.ib(
+    data: Dict[str, Union[str, pathlib.Path, ColorMapType]] = attr.ib(
         default=attr.Factory(lambda: DEFAULT_CMAPS_FILES)
     )
 
@@ -288,17 +296,45 @@ class ColorMaps:
             dict: colormap dictionary.
 
         """
-        cmap = self.data.get(name.lower(), None)
+        cmap = self.data.get(name, None)
         if cmap is None:
             raise InvalidColorMapName(f"Invalid colormap name: {name}")
 
-        if isinstance(cmap, str):
-            colormap = numpy.load(cmap)
-            assert colormap.shape == (256, 4)
-            assert colormap.dtype == numpy.uint8
-            return {idx: tuple(value) for idx, value in enumerate(colormap)}  # type: ignore
-        else:
-            return cmap
+        if isinstance(cmap, (pathlib.Path, str)):
+            if isinstance(cmap, str):
+                cmap = pathlib.Path(cmap)
+
+            if cmap.suffix == ".npy":
+                colormap = numpy.load(cmap)
+                assert colormap.shape == (256, 4)
+                assert colormap.dtype == numpy.uint8
+                cmap_data = {idx: tuple(value) for idx, value in enumerate(colormap)}
+
+            elif cmap.suffix == ".json":
+                with cmap.open() as f:
+                    cmap_data = json.load(
+                        f,
+                        object_hook=lambda x: {
+                            int(k): parse_color(v) for k, v in x.items()
+                        },
+                    )
+
+                # Make sure to match colormap type
+                if isinstance(cmap_data, Sequence):
+                    cmap_data = [
+                        (tuple(inter), parse_color(v))  # type: ignore
+                        for (inter, v) in cmap_data
+                    ]
+
+            else:
+                raise ValueError(f"Not supported {cmap.suffix} extension for ColorMap")
+
+            # save the numpy array / dict / sequence in the data dict
+            # avoiding the need to re-load the data
+            self.data[name] = cmap_data
+            return cmap_data
+
+        return cmap
 
     def list(self) -> List[str]:
         """List registered Colormaps.
@@ -311,7 +347,7 @@ class ColorMaps:
 
     def register(
         self,
-        custom_cmap: Dict[str, Union[str, ColorMapType]],
+        custom_cmap: Dict[str, Union[str, pathlib.Path, ColorMapType]],
         overwrite: bool = False,
     ) -> "ColorMaps":
         """Register a custom colormap.