fix: sc.pl.paga cax handling with multiple colors

docs/release-notes/4319.fix.md

@@ -0,0 +1 @@
+Fix `sc.pl.paga` `TypeError` when a single `cax` was passed together with multiple `color`s {smaller}`Advait Shukla`

src/scanpy/plotting/legacy/_tools/paga.py

@@ -381,7 +381,7 @@ def paga(  # noqa: PLR0912, PLR0913, PLR0915
     pos: np.ndarray | Path | str | None = None,
     normalize_to_color: bool = False,
     cmap: str | Colormap | None = None,
-    cax: Axes | None = None,
+    cax: Axes | Sequence[Axes] | None = None,
     colorbar=None,  # TODO: this seems to be unused
     cb_kwds: Mapping[str, Any] = frozendict({}),
     frameon: bool | None = None,
@@ -491,7 +491,8 @@ def paga(  # noqa: PLR0912, PLR0913, PLR0915
     cmap
         The color map.
     cax
-        A matplotlib axes object for a potential colorbar.
+        A matplotlib axes object, or a sequence of axes (one per color), for
+        a potential colorbar.
     cb_kwds
         Keyword arguments for :class:`~matplotlib.colorbar.Colorbar`,
         for instance, `ticks`.
@@ -685,8 +686,16 @@ def is_flat(x):
                     rectangle = [left, bottom, width, height]
                     fig = plt.gcf()
                     ax_cb = fig.add_axes(rectangle)
-                else:
+                elif isinstance(cax, (list, tuple, np.ndarray)):
                     ax_cb = cax[icolor]
+                else:
+                    if sum(colorbars) > 1:
+                        msg = (
+                            "`cax` must be a sequence of axes (one per color) "
+                            "when multiple colorbars are requested."
+                        )
+                        raise ValueError(msg)
+                    ax_cb = cax
 
                 _ = plt.colorbar(
                     sct,

tests/plotting/legacy/test_paga.py

@@ -3,9 +3,13 @@
 from functools import partial
 from importlib.util import find_spec
 
+import numpy as np
+import pandas as pd
 import pytest
 from matplotlib import colormaps
+from matplotlib import pyplot as plt
 from packaging.version import Version
+from scipy import sparse
 
 import scanpy as sc
 from scanpy._compat import pkg_version
@@ -96,3 +100,34 @@ def test_paga_compare(plot_cmp):
     sc.pl.paga_compare(pbmc, basis="umap", show=False)
 
     plot_cmp("paga_compare_pbmc3k")
+
+
+def test_paga_cax() -> None:
+    # Tests that https://github.com/scverse/scanpy/issues/4318 is fixed
+    rng = np.random.default_rng(0)
+    adata = sc.AnnData(rng.random((80, 20)))
+    adata.obs["group"] = pd.Categorical(rng.choice(["a", "b", "c", "d", "e"], 80))
+
+    k = 5
+    rows = np.array([0, 1, 1, 2, 2, 3, 3, 4, 4, 0])
+    cols = np.array([1, 0, 2, 1, 3, 2, 4, 3, 0, 4])
+    connectivities = sparse.csr_matrix(  # noqa: TID251
+        (np.ones(len(rows)), (rows, cols)), shape=(k, k)
+    )
+    adata.uns["paga"] = {
+        "groups": "group",
+        "connectivities": connectivities,
+        "connectivities_tree": connectivities.copy(),
+    }
+    pos = rng.random((k, 2))
+
+    # a single `cax` works for a single colorbar
+    _, cax = plt.subplots()
+    sc.pl.paga(adata, color=adata.var_names[0], cax=cax, pos=pos, show=False)
+
+    # a single `cax` with multiple colorbars raises a clear error
+    _, cax = plt.subplots()
+    with pytest.raises(ValueError, match="sequence of axes"):
+        sc.pl.paga(
+            adata, color=adata.var_names[:2].tolist(), cax=cax, pos=pos, show=False
+        )