Skip to content
Back to top

# Source code for PIL.ImagePalette

    #
    # The Python Imaging Library.
    # $Id$
    #
    # image palette object
    #
    # History:
    # 1996-03-11 fl   Rewritten.
    # 1997-01-03 fl   Up and running.
    # 1997-08-23 fl   Added load hack
    # 2001-04-16 fl   Fixed randint shadow bug in random()
    #
    # Copyright (c) 1997-2001 by Secret Labs AB
    # Copyright (c) 1996-1997 by Fredrik Lundh
    #
    # See the README file for information on usage and redistribution.
    #
    from __future__ import annotations

    import array
    from collections.abc import Sequence
    from typing import IO

    from . import GimpGradientFile, GimpPaletteFile, ImageColor, PaletteFile

    TYPE_CHECKING = False
    if TYPE_CHECKING:
        from . import Image

    class ImagePalette:
        """
        Color palette for palette mapped images

        :param mode: The mode to use for the palette. See:
            :ref:`concept-modes`. Defaults to "RGB"
        :param palette: An optional palette. If given, it must be a bytearray,
            an array or a list of ints between 0-255. The list must consist of
            all channels for one color followed by the next color (e.g. RGBRGBRGB).
            Defaults to an empty palette.
        """

        def __init__(
            self,
            mode: str = "RGB",
            palette: Sequence[int] | bytes | bytearray | None = None,
        ) -> None:
            self.mode = mode
            self.rawmode: str | None = None  # if set, palette contains raw data
            self.palette = palette or bytearray()
            self.dirty: int | None = None

        @property
        def palette(self) -> Sequence[int] | bytes | bytearray:
            return self._palette

        @palette.setter
        def palette(self, palette: Sequence[int] | bytes | bytearray) -> None:
            self._colors: dict[tuple[int, ...], int] | None = None
            self._palette = palette

        @property
        def colors(self) -> dict[tuple[int, ...], int]:
            if self._colors is None:
                mode_len = len(self.mode)
                self._colors = {}
                for i in range(0, len(self.palette), mode_len):
                    color = tuple(self.palette[i : i + mode_len])
                    if color in self._colors:
                        continue
                    self._colors[color] = i // mode_len
            return self._colors

        @colors.setter
        def colors(self, colors: dict[tuple[int, ...], int]) -> None:
            self._colors = colors

        def copy(self) -> ImagePalette:
            new = ImagePalette()

            new.mode = self.mode
            new.rawmode = self.rawmode
            if self.palette is not None:
                new.palette = self.palette[:]
            new.dirty = self.dirty

            return new

        def getdata(self) -> tuple[str, Sequence[int] | bytes | bytearray]:
            """
            Get palette contents in format suitable for the low-level
            ``im.putpalette`` primitive.

            .. warning:: This method is experimental.
            """
            if self.rawmode:
                return self.rawmode, self.palette
            return self.mode, self.tobytes()

        def tobytes(self) -> bytes:
            """Convert palette to bytes.

            .. warning:: This method is experimental.
            """
            if self.rawmode:
                msg = "palette contains raw palette data"
                raise ValueError(msg)
            if isinstance(self.palette, bytes):
                return self.palette
            arr = array.array("B", self.palette)
            return arr.tobytes()

        # Declare tostring as an alias for tobytes
        tostring = tobytes

        def _new_color_index(
            self, image: Image.Image | None = None, e: Exception | None = None
        ) -> int:
            if not isinstance(self.palette, bytearray):
                self._palette = bytearray(self.palette)
            index = len(self.palette) // len(self.mode)
            special_colors: tuple[int | tuple[int, ...] | None, ...] = ()
            if image:
                special_colors = (
                    image.info.get("background"),
                    image.info.get("transparency"),
                )
                while index in special_colors:
                    index += 1
            if index >= 256:
                if image:
                    # Search for an unused index
                    for i, count in reversed(list(enumerate(image.histogram()))):
                        if count == 0 and i not in special_colors:
                            index = i
                            break
                if index >= 256:
                    msg = "cannot allocate more than 256 colors"
                    raise ValueError(msg) from e
            return index

        def getcolor(
            self,
            color: tuple[int, ...],
            image: Image.Image | None = None,
        ) -> int:
            """Given an rgb tuple, allocate palette entry.

            .. warning:: This method is experimental.
            """
            if self.rawmode:
                msg = "palette contains raw palette data"
                raise ValueError(msg)
            if isinstance(color, tuple):
                if self.mode == "RGB":
                    if len(color) == 4:
                        if color[3] != 255:
                            msg = "cannot add non-opaque RGBA color to RGB palette"
                            raise ValueError(msg)
                        color = color[:3]
                elif self.mode == "RGBA":
                    if len(color) == 3:
                        color += (255,)
                try:
                    return self.colors[color]
                except KeyError as e:
                    # allocate new color slot
                    index = self._new_color_index(image, e)
                    assert isinstance(self._palette, bytearray)
                    self.colors[color] = index
                    mode_len = len(self.mode)
                    if index * mode_len < len(self.palette):
                        self._palette = (
                            self._palette[: index * mode_len]
                            + bytes(color)
                            + self._palette[index * mode_len + mode_len :]
                        )
                    else:
                        self._palette += bytes(color)
                    self.dirty = 1
                    return index
            else:
                msg = f"unknown color specifier: {repr(color)}"  # type: ignore[unreachable]
                raise ValueError(msg)

        def save(self, fp: str | IO[str]) -> None:
            """Save palette to text file.

            .. warning:: This method is experimental.
            """
            if self.rawmode:
                msg = "palette contains raw palette data"
                raise ValueError(msg)
            open_fp = False
            if isinstance(fp, str):
                fp = open(fp, "w")
                open_fp = True
            try:
                fp.write("# Palette\n")
                fp.write(f"# Mode: {self.mode}\n")
                palette_len = len(self.palette)
                for i in range(256):
                    fp.write(f"{i}")
                    for j in range(i * len(self.mode), (i + 1) * len(self.mode)):
                        fp.write(f" {self.palette[j] if j < palette_len else 0}")
                    fp.write("\n")
            finally:
                if open_fp:
                    fp.close()

    # --------------------------------------------------------------------
    # Internal

    def raw(rawmode: str, data: Sequence[int] | bytes | bytearray) -> ImagePalette:
        palette = ImagePalette()
        palette.rawmode = rawmode
        palette.palette = data
        palette.dirty = 1
        return palette

    # --------------------------------------------------------------------
    # Factories

    def make_linear_lut(black: int, white: float) -> list[int]:
        if black == 0:
            return [int(white * i // 255) for i in range(256)]

        msg = "unavailable when black is non-zero"
        raise NotImplementedError(msg)  # FIXME

    def make_gamma_lut(exp: float) -> list[int]:
        return [int(((i / 255.0) ** exp) * 255.0 + 0.5) for i in range(256)]

    def negative(mode: str = "RGB") -> ImagePalette:
        palette = list(range(256 * len(mode)))
        palette.reverse()
        return ImagePalette(mode, [i // len(mode) for i in palette])

    def random(mode: str = "RGB") -> ImagePalette:
        from random import randint

        palette = [randint(0, 255) for _ in range(256 * len(mode))]
        return ImagePalette(mode, palette)

    def sepia(white: str = "#fff0c0") -> ImagePalette:
        bands = [make_linear_lut(0, band) for band in ImageColor.getrgb(white)]
        return ImagePalette("RGB", [bands[i % 3][i // 3] for i in range(256 * 3)])

    def wedge(mode: str = "RGB") -> ImagePalette:
        palette = list(range(256 * len(mode)))
        return ImagePalette(mode, [i // len(mode) for i in palette])

    def load(filename: str) -> tuple[bytes, str]:
        # FIXME: supports GIMP gradients only

        with open(filename, "rb") as fp:
            paletteHandlers: list[
                type[
                    GimpPaletteFile.GimpPaletteFile
                    | GimpGradientFile.GimpGradientFile
                    | PaletteFile.PaletteFile
                ]
            ] = [
                GimpPaletteFile.GimpPaletteFile,
                GimpGradientFile.GimpGradientFile,
                PaletteFile.PaletteFile,
            ]
            for paletteHandler in paletteHandlers:
                try:
                    fp.seek(0)
                    lut = paletteHandler(fp).getpalette()
                    if lut:
                        break
                except (SyntaxError, ValueError):
                    pass
            else:
                msg = "cannot load palette"
                raise OSError(msg)

        return lut  # data, rawmode