Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 22 additions & 8 deletions src/cmap/_colormap.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
from ._catalog import CatalogItem

LutCacheKey = tuple[int, float, bool]
RGBABytes = tuple[np.uint8, np.uint8, np.uint8, np.uint8]
Interpolation = Literal["linear", "nearest"]

class MPLSegmentData(TypedDict, total=False):
Expand Down Expand Up @@ -407,16 +408,24 @@ def __call__(
*,
N: int = 256,
gamma: float = 1,
bytes: bool = False,
bytes: Literal[False] = False,
) -> Color: ...
@overload
def __call__(
self, x: float, *, N: int = 256, gamma: float = 1, bytes: Literal[True]
) -> RGBABytes: ...
@overload
def __call__(
self, x: float, *, N: int = 256, gamma: float = 1, bytes: bool
) -> Color | RGBABytes: ...
def __call__(
self,
x: float | NDArray | Iterable[float],
*,
N: int = 256,
gamma: float = 1,
bytes: bool = False,
) -> Color | NDArray[np.float64]:
) -> Color | RGBABytes | NDArray[np.float64]:
r"""Map scalar values in X to an RGBA array.

This is the primary API for "using" a `cmap.Colormap` to map scalar values to
Expand Down Expand Up @@ -460,14 +469,15 @@ def __call__(
bytes : bool
If False (default), the returned RGBA values will be floats in the
interval ``[0, 1]`` otherwise they will be `numpy.uint8`\s in the
interval ``[0, 255]``.
interval ``[0, 255]``. For scalar `x`, `bytes=True` returns a tuple of
four `numpy.uint8` (as matplotlib does) rather than a `Color`.

Returns
-------
color : Color | NDArray
If X is a float, a single RGBA color will be returned. If x is an
array-like, an array of RGBA colors will be returned with shape
`x.shape + (4,)`
color : Color | tuple | NDArray
If X is a float, a single `Color` will be returned (or a 4-tuple of
`numpy.uint8` if `bytes` is True). If x is an array-like, an array of
RGBA colors will be returned with shape `x.shape + (4,)`

Examples
--------
Expand Down Expand Up @@ -532,7 +542,11 @@ def __call__(
xa[mask_masked] = N + 2

rgba = lut.take(xa, axis=0, mode="clip")
return rgba if np.iterable(x) else Color(rgba)
if np.iterable(x):
return rgba
if bytes:
return tuple(rgba)
return Color(rgba)

def _with_exceptional_colors(self, lut: np.ndarray) -> np.ndarray:
"""Return `lut` with three rows appended, one per exceptional value class.
Expand Down
11 changes: 11 additions & 0 deletions tests/test_colormap.py
Original file line number Diff line number Diff line change
Expand Up @@ -591,3 +591,14 @@ def test_invalid_function_colormap_with_cmap_kwargs() -> None:
# providing cmap_kwargs to a non-callable colormap should raise TypeError
with pytest.raises(TypeError, match=r"Cannot apply cmap_kwargs to colormap"):
Colormap("viridis", cmap_kwargs={"param": 1.0})


def test_scalar_bytes() -> None:
cm = Colormap("viridis")
expect = cm(np.array([0.5]), bytes=True)[0]
for x in (0.5, np.float64(0.5), np.array(0.5)):
rgba = cm(x, bytes=True)
assert isinstance(rgba, tuple)
assert rgba == tuple(expect)
assert all(isinstance(v, np.uint8) for v in rgba)
assert isinstance(cm(0.5), Color)
Loading