(4.3.0) Add strict type checking using pyright and improve the static types (#85)

* Replace mypy with pyright

* Add a types.py file containing custom types

* Remove Union in favor of |, replace pre-commit hooks with local tools

* Fix an issue with np.argmax and ColorBox not implementing the __array__ protocol

* Bump minor version, add changelog

* Remove redundant comment

Author: qTipTip

.github/workflows/ci_tests.yaml

@@ -29,5 +29,9 @@ jobs:
           key: venv-${{ hashFiles('uv.lock') }}
       - name: Install the project dependencies
         run: uv sync --extra dev
+      - name: Install pyright
+        run: uv tool install pyright
+      - name: Run type checking
+        run: uv run pyright
       - name: Run the automated tests (for example)
         run: uv run pytest -v

.pre-commit-config.yaml

@@ -6,27 +6,23 @@ repos:
         exclude: ^mkdocs\.yml$
       - id: end-of-file-fixer
       - id: trailing-whitespace
-  - repo: https://github.com/astral-sh/ruff-pre-commit
-    # Ruff version.
-    rev: v0.5.0
+  - repo: local
     hooks:
-      # Run the linter.
-      - id: ruff
-        args:
-          - --config=pyproject.toml
-          - --fix
-      # Run the formatter.
+      - id: ruff-check
+        name: ruff check
+        entry: uv run ruff check --config=pyproject.toml --fix
+        language: system
+        types: [python]
+        require_serial: false
       - id: ruff-format
-  -   repo: https://github.com/pre-commit/mirrors-mypy
-      rev: 'v1.10.0'  # Use the sha / tag you want to point at
-      hooks:
-      -   id: mypy
-          args:
-            - --show-error-codes
-            - --pretty
-            - --warn-redundant-casts
-            - --check-untyped-defs
-            - --ignore-missing-imports
-            - --disallow-any-generics
-            - --disallow-subclassing-any
-            - --no-implicit-optional
+        name: ruff format
+        entry: uv run ruff format
+        language: system
+        types: [python]
+        require_serial: false
+      - id: pyright
+        name: pyright
+        entry: uv run pyright
+        language: system
+        types: [python]
+        require_serial: false

CHANGELOG.md

@@ -10,13 +10,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 # Released
 
-## 4.2.0 - 08/11/2025
+## 4.3.0 - 11/08/2025
+
+### Added
+
+• Added `types.py` file containing custom type definitions
+
+### Changed
+
+• Replaced mypy with pyright for type checking
+• Updated and improved type annotations
+• Replaced pre-commit hooks with local tools
+
+### Fixed
+
+• Fixed an issue with np.argmax and ColorBox not implementing the __array__ protocol, improving NumPy compatibility
+
+## 4.2.0 - 11/08/2025
 
 ### Added
 
 - Added support for `PIL.Image` type in `extract_colors`.
 
-## 4.1.0 - 4/7/2025
+## 4.1.0 - 04/07/2025
 
 ### Added
 

Pylette/__init__.py

@@ -1,5 +1,6 @@
 from Pylette.src.color import Color
-from Pylette.src.color_extraction import ImageType_T, extract_colors
+from Pylette.src.color_extraction import extract_colors
 from Pylette.src.palette import Palette
+from Pylette.src.types import ImageInput
 
-__all__ = ["extract_colors", "Palette", "Color", "ImageType_T"]
+__all__ = ["extract_colors", "Palette", "Color", "ImageInput"]

Pylette/cmd.py

@@ -52,11 +52,11 @@ def main(
         typer.echo("Please provide either a filename or an image-url, but not both.")
         raise typer.Exit(code=1)
 
-    image: pathlib.Path | str | None
-    if filename is not None and image_url is None:
-        image = filename
+    if filename is not None:
+        image = filename  # Path
     else:
-        image = image_url
+        assert image_url is not None
+        image = image_url  # str (URL)
 
     output_file_path = str(out_filename) if out_filename is not None else None
     palette = extract_colors(

Pylette/src/color_extraction.py

@@ -1,68 +1,50 @@
-import os
 import urllib.parse
-from enum import Enum
 from io import BytesIO
-from typing import Any, Literal, TypeAlias, Union
+from pathlib import Path
+from typing import Literal
 
 import numpy as np
-import requests  # type: ignore
-from numpy.typing import NDArray
+import requests
 from PIL import Image
 
 from Pylette.src.extractors.k_means import k_means_extraction
 from Pylette.src.extractors.median_cut import median_cut_extraction
 from Pylette.src.palette import Palette
-
-ImageType_T: TypeAlias = Union["os.PathLike[Any]", bytes, NDArray[float], str, Image.Image]
-
-
-class ImageType(str, Enum):
-    PATH = "path"
-    BYTES = "bytes"
-    ARRAY = "array"
-    URL = "url"
-    PIL = "pil"
-    NONE = "none"
-
-
-def _parse_image_type(image: ImageType_T) -> ImageType:
-    """
-    Determines the type of the input image.
-
-    Parameters:
-    image (ImageType_T): The input image.
-
-    Returns:
-    ImageType: The type of the input image.
-    """
-    match image:
-        case Image.Image():
-            image_type = ImageType.PIL
-        case np.ndarray():
-            image_type = ImageType.ARRAY
-        case os.PathLike():
-            image_type = ImageType.PATH
-        case bytes():
-            image_type = ImageType.BYTES
-        case str():
-            try:
-                result = urllib.parse.urlparse(image)
-                if all([result.scheme, result.netloc]):
-                    image_type = ImageType.URL
-                else:
-                    image_type = ImageType.PATH
-            except ValueError:
-                image_type = ImageType.PATH
-        case _:
-            image_type = ImageType.NONE
-    return image_type
+from Pylette.src.types import ImageInput, PILImage
+
+
+def _is_url(image_str: str) -> bool:
+    """Check if a string is a valid URL."""
+    try:
+        result = urllib.parse.urlparse(image_str)
+        return all([result.scheme, result.netloc])
+    except Exception:
+        return False
+
+
+def _normalize_image_input(image: ImageInput) -> PILImage:
+    """Convert any valid image input to PIL Image."""
+    if isinstance(image, Image.Image):
+        return image
+    elif isinstance(image, (str, Path)):
+        image_str = str(image)
+        if _is_url(image_str):
+            return request_image(image_str)
+        else:
+            return Image.open(image)
+    elif isinstance(image, bytes):
+        return Image.open(BytesIO(image))
+    elif hasattr(image, "__array__"):  # More general check for array-like objects
+        return Image.fromarray(image)
+    else:
+        raise TypeError(f"Unsupported image type: {type(image)}")
 
 
 def extract_colors(
-    image: ImageType_T,
+    image: ImageInput,
     palette_size: int = 5,
     resize: bool = True,
-    mode: Literal["KM"] | Literal["MC"] = "KM",
+    mode: Literal["KM", "MC"] = "KM",
     sort_mode: Literal["luminance", "frequency"] | None = None,
     alpha_mask_threshold: int | None = None,
 ) -> Palette:
@@ -87,27 +69,9 @@ def extract_colors(
         >>> extract_colors(b"image_bytes", palette_size=5, resize=True, mode="KM", sort_mode="luminance")
     """
 
-    image_type = _parse_image_type(image)
-
-    match image_type:
-        case ImageType.PATH:
-            img_obj = Image.open(image)
-        case ImageType.BYTES:
-            assert isinstance(image, bytes)
-            img_obj = Image.open(BytesIO(image))
-        case ImageType.URL:
-            assert isinstance(image, str)
-            img_obj = request_image(image)
-        case ImageType.ARRAY:
-            img_obj = Image.fromarray(image)
-        case ImageType.PIL:
-            img_obj = image
-        case ImageType.NONE:
-            raise ValueError(f"Unable to parse image source. Got image type {type(image)}")
-
-    # Convert to RGBA
+    # Normalize input to PIL Image and convert to RGBA
+    img_obj = _normalize_image_input(image)
     img = img_obj.convert("RGBA")
-
     # open the image
     if resize:
         img = img.resize((256, 256))
@@ -121,12 +85,11 @@ def extract_colors(
     alpha_mask = arr[:, :, 3] <= alpha_mask_threshold
     valid_pixels = arr[~alpha_mask]
 
-    if mode == "KM":
-        colors = k_means_extraction(valid_pixels, height, width, palette_size)
-    elif mode == "MC":
-        colors = median_cut_extraction(valid_pixels, height, width, palette_size)
-    else:
-        raise NotImplementedError("Extraction mode not implemented")
+    match mode:
+        case "KM":
+            colors = k_means_extraction(valid_pixels, height, width, palette_size)
+        case "MC":
+            colors = median_cut_extraction(valid_pixels, height, width, palette_size)
 
     if sort_mode == "luminance":
         colors.sort(key=lambda c: c.luminance, reverse=False)

Pylette/src/extractors/median_cut.py

@@ -25,9 +25,8 @@ def extract(self, arr: NDArray[NP_T], height: int, width: int, palette_size: int
         valid_pixel_count = arr.shape[0]
         boxes = [ColorBox(arr)]
         while len(boxes) < palette_size:
-            largest_box_idx = np.argmax(boxes)  # type: ignore
+            largest_box_idx = np.argmax([box.size for box in boxes])
             boxes = boxes[:largest_box_idx] + boxes[largest_box_idx].split() + boxes[largest_box_idx + 1 :]
-
         return [Color(tuple(map(int, box.average)), box.pixel_count / valid_pixel_count) for box in boxes]
 
 

Pylette/src/palette.py

@@ -115,7 +115,7 @@ def to_csv(
                         palette_file.write(",{}".format(color.freq))
                     palette_file.write("\n")
 
-    def random_color(self, N, mode="frequency"):
+    def random_color(self, N: int, mode: str = "frequency") -> list[Color]:
         """
         Returns N random colors from the palette, either using the frequency of each color, or choosing uniformly.
 
@@ -128,11 +128,17 @@ def random_color(self, N, mode="frequency"):
         """
 
         if mode == "frequency":
-            pdf = self.frequencies
+            # Convert to numpy-compatible format for weighted selection
+            colors_array = np.array(range(len(self.colors)))
+            indices = np.random.choice(colors_array, size=N, p=self.frequencies)
+            return [self.colors[i] for i in indices]
         elif mode == "uniform":
-            pdf = None
-
-        return np.random.choice(self.colors, size=N, p=pdf)
+            # Uniform selection without weights
+            colors_array = np.array(range(len(self.colors)))
+            indices = np.random.choice(colors_array, size=N)
+            return [self.colors[i] for i in indices]
+        else:
+            raise ValueError(f"Invalid mode: {mode}. Must be 'frequency' or 'uniform'.")
 
     def __str__(self):
         return "".join(["({}, {}, {}, {}) \n".format(c.rgb[0], c.rgb[1], c.rgb[2], c.freq) for c in self.colors])

Pylette/src/types.py

@@ -0,0 +1,46 @@
+"""
+Centralized type definitions for Pylette.
+
+This module contains all the type aliases and protocols used throughout the Pylette library
+to ensure type safety and consistency.
+"""
+
+from pathlib import Path
+from typing import Any, Protocol, TypeAlias
+
+import numpy as np
+from numpy.typing import NDArray
+from PIL import Image
+
+
+class ImageLike(Protocol):
+    """Protocol for image-like objects that can be converted to PIL Image."""
+
+    pass
+
+
+class ArrayLike(Protocol):
+    """Protocol for array-like objects."""
+
+    def __array__(self) -> NDArray[np.uint8]: ...
+
+
+# Specific image input types
+PathLikeImage: TypeAlias = str | Path
+URLImage: TypeAlias = str  # URLs are strings but semantically different
+BytesImage: TypeAlias = bytes
+ArrayImage: TypeAlias = NDArray[np.uint8]  # Properly typed array
+PILImage: TypeAlias = Image.Image
+
+# Main union type - more restrictive and logical
+ImageInput: TypeAlias = PathLikeImage | URLImage | BytesImage | ArrayImage | PILImage
+
+# Color array types
+ColorArray: TypeAlias = NDArray[np.uint8]  # For RGB/RGBA color data
+FloatArray: TypeAlias = NDArray[np.floating[Any]]  # For calculations
+IntArray: TypeAlias = NDArray[np.integer[Any]]  # For integer arrays
+
+# Color tuple types
+RGBTuple: TypeAlias = tuple[int, int, int]
+RGBATuple: TypeAlias = tuple[int, int, int, int]
+ColorTuple: TypeAlias = RGBTuple | RGBATuple