Patch/better self typing def (#736)

* improve type hints

* improve type hints

.pre-commit-config.yaml

@@ -18,7 +18,7 @@ repos:
       - id: ruff-format
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.9.0
+    rev: v1.11.2
     hooks:
       - id: mypy
         language_version: python

CHANGES.md

@@ -10,6 +10,8 @@
 * Moved `_dst_geom_in_tms_crs` from Reader to `SpatialMixin` class **breaking change**
 * Removed use of rasterio's `is_tiled` method
 * Enable **Alternate** asset's HREF for STAC by using `RIO_TILER_STAC_ALTERNATE_KEY` environment variable
+* Improve type hint definition
+* make `ImageData.rescale` and `ImageData.apply_color_formula` to return `self`
 
 # 6.7.0 (2024-09-05)
 

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
     "pystac>=0.5.4",
     "rasterio>=1.3.0",
     "color-operations",
+    "typing-extensions",
     "importlib_resources>=1.1.0; python_version < '3.9'",
 ]
 

rio_tiler/io/base.py

@@ -346,7 +346,9 @@ class MultiBaseReader(SpatialMixin, metaclass=abc.ABCMeta):
     assets: Sequence[str] = attr.ib(init=False)
     default_assets: Optional[Sequence[str]] = attr.ib(init=False, default=None)
 
-    ctx: Any = attr.ib(init=False, default=contextlib.nullcontext)
+    ctx: Type[contextlib.AbstractContextManager] = attr.ib(
+        init=False, default=contextlib.nullcontext
+    )
 
     def __enter__(self):
         """Support using with Context Managers."""

rio_tiler/io/rasterio.py

@@ -82,7 +82,7 @@ class Reader(BaseReader):
     options: reader.Options = attr.ib()
 
     # Context Manager to handle rasterio open/close
-    _ctx_stack = attr.ib(init=False, factory=contextlib.ExitStack)
+    _ctx_stack: contextlib.ExitStack = attr.ib(init=False, factory=contextlib.ExitStack)
 
     @options.default
     def _options_default(self):

rio_tiler/io/stac.py

@@ -250,7 +250,7 @@ class STACReader(MultiBaseReader):
 
     fetch_options: Dict = attr.ib(factory=dict)
 
-    ctx: Any = attr.ib(default=rasterio.Env)
+    ctx: rasterio.Env = attr.ib(default=rasterio.Env)
 
     def __attrs_post_init__(self):
         """Fetch STAC Item and get list of valid assets."""

rio_tiler/models.py

@@ -21,6 +21,7 @@
 from rasterio.plot import reshape_as_image
 from rasterio.transform import from_bounds
 from rasterio.warp import transform_geom
+from typing_extensions import Self
 
 from rio_tiler.colormap import apply_cmap
 from rio_tiler.constants import WGS84_CRS
@@ -218,7 +219,7 @@ def count(self) -> int:
         return self.array.shape[0]
 
     @classmethod
-    def create_from_list(cls, data: Sequence["PointData"]):
+    def create_from_list(cls, data: Sequence["PointData"]) -> Self:
         """Create PointData from a sequence of PointsData objects.
 
         Args:
@@ -361,7 +362,7 @@ def __iter__(self):
             yield i
 
     @classmethod
-    def from_array(cls, arr: numpy.ndarray) -> "ImageData":
+    def from_array(cls, arr: numpy.ndarray) -> Self:
         """Create ImageData from a numpy array.
 
         Args:
@@ -376,7 +377,7 @@ def from_array(cls, arr: numpy.ndarray) -> "ImageData":
         return cls(arr)
 
     @classmethod
-    def from_bytes(cls, data: bytes) -> "ImageData":
+    def from_bytes(cls, data: bytes) -> Self:
         """Create ImageData from bytes.
 
         Args:
@@ -425,7 +426,7 @@ def from_bytes(cls, data: bytes) -> "ImageData":
                     )
 
     @classmethod
-    def create_from_list(cls, data: Sequence["ImageData"]) -> "ImageData":
+    def create_from_list(cls, data: Sequence["ImageData"]) -> Self:
         """Create ImageData from a sequence of ImageData objects.
 
         Args:
@@ -531,7 +532,7 @@ def count(self) -> int:
         return self.array.shape[0]
 
     @property
-    def transform(self):
+    def transform(self) -> Affine:
         """Returns the affine transform."""
         return (
             from_bounds(*self.bounds, self.width, self.height)
@@ -544,14 +545,15 @@ def rescale(
         in_range: Sequence[IntervalTuple],
         out_range: Sequence[IntervalTuple] = ((0, 255),),
         out_dtype: Union[str, numpy.number] = "uint8",
-    ):
+    ) -> Self:
         """Rescale data in place."""
         self.array = rescale_image(
             self.array.copy(),
             in_range=in_range,
             out_range=out_range,
             out_dtype=out_dtype,
         )
+        return self
 
     def apply_colormap(self, colormap: ColorMapType) -> "ImageData":
         """Apply colormap to the image data."""
@@ -570,7 +572,7 @@ def apply_colormap(self, colormap: ColorMapType) -> "ImageData":
             metadata=self.metadata,
         )
 
-    def apply_color_formula(self, color_formula: Optional[str]):
+    def apply_color_formula(self, color_formula: Optional[str]) -> Self:
         """Apply color-operations formula in place."""
         out = self.array.data.copy()
         out[out < 0] = 0
@@ -581,6 +583,7 @@ def apply_color_formula(self, color_formula: Optional[str]):
         data = numpy.ma.MaskedArray(out)
         data.mask = self.array.mask
         self.array = data
+        return self
 
     def apply_expression(self, expression: str) -> "ImageData":
         """Apply expression to the image data."""

rio_tiler/utils.py

@@ -397,7 +397,7 @@ def get_vrt_transform(
         output_transform = from_bounds(w, s, e, n, width, height)
 
         # NOTE: Here we check if the Output Resolution is higher thant the dataset resolution (OverZoom)
-        # When not overzooming we don't want to use the output Width/Height to calculate the transform
+        # When not over-zooming we don't want to use the output Width/Height to calculate the transform
         # See issues https://github.com/cogeotiff/rio-tiler/pull/648
         w_res = (
             output_transform.a