catppuccin.extras.matplotlib

  1"""Soothing pastel theme for `matplotlib`.
  2
  3The following code was ported from [catppuccin/matplotlib](https://github.com/catppuccin/matplotlib).
  4Thanks to [Bram de Wilde](https://github.com/brambozz) for the original source code and
  5for allowing this port.
  6
  7The library tries to register styles and colormaps if `matplotlib` is installed.
  8See the examples below for some use cases:
  9
 101. Load a style, using `mpl.style.use`
 11
 12   ```python
 13   import catppuccin
 14   import matplotlib as mpl
 15   import matplotlib.pyplot as plt
 16
 17   mpl.style.use(catppuccin.PALETTE.mocha.identifier)
 18   plt.plot([0,1,2,3], [1,2,3,4])
 19   plt.show()
 20   ```
 21
 221. Mix it with different stylesheets!
 23
 24   ```python
 25   import catppuccin
 26   import matplotlib as mpl
 27   import matplotlib.pyplot as plt
 28
 29   mpl.style.use(["ggplot", catppuccin.PALETTE.mocha.identifier])
 30   plt.plot([0,1,2,3], [1,2,3,4])
 31   plt.show()
 32   ```
 33
 341. Load individual colors
 35
 36   ```python
 37   import matplotlib.pyplot as plt
 38   import catppuccin
 39   from catppuccin.extras.matplotlib import load_color
 40
 41   color = load_color(catppuccin.PALETTE.latte.identifier, "peach")
 42   plt.plot([0,1,2,3], [1,2,3,4], color=color)
 43   plt.show()
 44   ```
 45
 461. Define custom colormaps
 47
 48   ```python
 49   import matplotlib.pyplot as plt
 50   import numpy as np
 51   import catppuccin
 52   from catppuccin.extras.matplotlib import get_colormap_from_list
 53
 54   cmap = get_colormap_from_list(
 55       catppuccin.PALETTE.frappe.identifier,
 56       ["red", "peach", "yellow", "green"],
 57   )
 58   rng = np.random.default_rng()
 59   data = rng.integers(2, size=(30, 30))
 60
 61   plt.imshow(data, cmap=cmap)
 62   plt.show()
 63   ```
 64"""
 65
 66from __future__ import annotations
 67
 68from dataclasses import asdict
 69from pathlib import Path
 70from typing import TYPE_CHECKING, cast
 71
 72import matplotlib as mpl
 73import matplotlib.colors
 74import matplotlib.style
 75
 76from catppuccin.palette import PALETTE
 77
 78if TYPE_CHECKING:
 79    from collections.abc import Iterable
 80
 81CATPPUCCIN_STYLE_DIRECTORY = Path(__file__).parent / "matplotlib_styles"
 82DEFAULT_COLORMAP_COLORS = ("base", "blue")
 83
 84
 85def _register_styles() -> None:
 86    """Register the included stylesheets in the mpl style library."""
 87    catppuccin_stylesheets = mpl.style.core.read_style_directory(  # type: ignore [attr-defined]
 88        CATPPUCCIN_STYLE_DIRECTORY
 89    )
 90    mpl.style.core.update_nested_dict(mpl.style.library, catppuccin_stylesheets)  # type: ignore [attr-defined]
 91
 92
 93def _register_colormap_list() -> None:
 94    """Register the included color maps in the `matplotlib` colormap library."""
 95    for palette_name in asdict(PALETTE):
 96        cmap = get_colormap_from_list(palette_name, DEFAULT_COLORMAP_COLORS)
 97        mpl.colormaps.register(cmap=cmap, name=palette_name, force=True)
 98
 99
100def get_colormap_from_list(
101    palette_name: str,
102    color_names: Iterable[str],
103) -> matplotlib.colors.LinearSegmentedColormap:
104    """Get a `matplotlib` colormap from a list of colors for a specific palette."""
105    colors = [load_color(palette_name, color_name) for color_name in color_names]
106    return matplotlib.colors.LinearSegmentedColormap.from_list(palette_name, colors)
107
108
109def load_color(palette_name: str, color_name: str) -> str:
110    """Load the color for a given palette and color name."""
111    return cast(
112        str,
113        PALETTE.__getattribute__(palette_name).colors.__getattribute__(color_name).hex,
114    )