Allow all valid legend_loc options (scverse#3163)

Author: flying-sheep

docs/release-notes/1.10.3.md

@@ -5,7 +5,7 @@
 import collections.abc as cabc
 from collections import OrderedDict
 from itertools import product
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, get_args
 
 import matplotlib as mpl
 import numpy as np
@@ -47,30 +47,14 @@
     from seaborn import FacetGrid
     from seaborn.matrix import ClusterGrid
 
-    from ._utils import ColorLike, _FontSize, _FontWeight
+    from ._utils import ColorLike, _FontSize, _FontWeight, _LegendLoc
 
     # TODO: is that all?
     _Basis = Literal["pca", "tsne", "umap", "diffmap", "draw_graph_fr"]
     _VarNames = Union[str, Sequence[str]]
 
 
-VALID_LEGENDLOCS = {
-    "none",
-    "right margin",
-    "on data",
-    "on data export",
-    "best",
-    "upper right",
-    "upper left",
-    "lower left",
-    "lower right",
-    "right",
-    "center left",
-    "center right",
-    "lower center",
-    "upper center",
-    "center",
-}
+VALID_LEGENDLOCS = frozenset(get_args(_utils._LegendLoc))
 
 
 @old_positionals(
@@ -105,7 +89,7 @@ def scatter(
     groups: str | Iterable[str] | None = None,
     components: str | Collection[str] | None = None,
     projection: Literal["2d", "3d"] = "2d",
-    legend_loc: str = "right margin",
+    legend_loc: _LegendLoc | None = "right margin",
     legend_fontsize: int | float | _FontSize | None = None,
     legend_fontweight: int | _FontWeight | None = None,
     legend_fontoutline: float | None = None,
@@ -202,7 +186,7 @@ def _scatter_obs(
     groups=None,
     components=None,
     projection: Literal["2d", "3d"] = "2d",
-    legend_loc="right margin",
+    legend_loc: _LegendLoc | None = "right margin",
     legend_fontsize=None,
     legend_fontweight=None,
     legend_fontoutline=None,
@@ -211,7 +195,7 @@ def _scatter_obs(
     frameon=None,
     right_margin=None,
     left_margin=None,
-    size=None,
+    size: int | float | None = None,
     marker=".",
     title=None,
     show=None,
@@ -379,7 +363,7 @@ def _scatter_obs(
             key.replace("_", " ") if not is_color_like(key) else "" for key in keys
         ]
 
-    axs = scatter_base(
+    axs: list[Axes] = scatter_base(
         Y,
         title=title,
         alpha=alpha,

docs/release-notes/1.9.7.md

@@ -80,8 +80,8 @@
 projection
     Projection of plot (default: `'2d'`).
 legend_loc
-    Location of legend, either `'on data'`, `'right margin'` or a valid keyword
-    for the `loc` parameter of :class:`~matplotlib.legend.Legend`.
+    Location of legend, either `'on data'`, `'right margin'`, `None`,
+    or a valid keyword for the `loc` parameter of :class:`~matplotlib.legend.Legend`.
 legend_fontsize
     Numeric size in pt or string describing the size.
     See :meth:`~matplotlib.text.Text.set_fontsize`.

docs/release-notes/1.9.8.md

@@ -31,7 +31,7 @@
     from matplotlib.axes import Axes
     from matplotlib.colors import Colormap
 
-    from .._utils import _FontSize, _FontWeight, _IGraphLayout
+    from .._utils import _FontSize, _FontWeight, _IGraphLayout, _LegendLoc
 
 
 @old_positionals(
@@ -67,7 +67,7 @@ def paga_compare(
     groups=None,
     components=None,
     projection: Literal["2d", "3d"] = "2d",
-    legend_loc="on data",
+    legend_loc: _LegendLoc | None = "on data",
     legend_fontsize: int | float | _FontSize | None = None,
     legend_fontweight: int | _FontWeight = "bold",
     legend_fontoutline=None,

src/scanpy/plotting/_anndata.py

@@ -52,6 +52,7 @@
     _FontSize,  # noqa: TCH001
     _FontWeight,  # noqa: TCH001
     _IGraphLayout,  # noqa: TCH001
+    _LegendLoc,  # noqa: TCH001
     check_colornorm,
     check_projection,
     circles,
@@ -97,7 +98,7 @@ def embedding(
     frameon: bool | None = None,
     legend_fontsize: int | float | _FontSize | None = None,
     legend_fontweight: int | _FontWeight = "bold",
-    legend_loc: str = "right margin",
+    legend_loc: _LegendLoc | None = "right margin",
     legend_fontoutline: int | None = None,
     colorbar_loc: str | None = "right",
     vmax: VBound | Sequence[VBound] | None = None,
@@ -1091,11 +1092,11 @@ def _components_to_dimensions(
 
 
 def _add_categorical_legend(
-    ax,
+    ax: Axes,
     color_source_vector,
     *,
     palette: dict,
-    legend_loc: str,
+    legend_loc: _LegendLoc | None,
     legend_fontweight,
     legend_fontsize,
     legend_fontoutline,
@@ -1124,17 +1125,7 @@ def _add_categorical_legend(
         box = ax.get_position()
         ax.set_position([box.x0, box.y0, box.width * 0.91, box.height])
 
-    if legend_loc == "right margin":
-        for label in cats:
-            ax.scatter([], [], c=palette[label], label=label)
-        ax.legend(
-            frameon=False,
-            loc="center left",
-            bbox_to_anchor=(1, 0.5),
-            ncol=(1 if len(cats) <= 14 else 2 if len(cats) <= 30 else 3),
-            fontsize=legend_fontsize,
-        )
-    elif legend_loc == "on data":
+    if legend_loc == "on data":
         # identify centroids to put labels
 
         all_pos = (
@@ -1158,6 +1149,19 @@ def _add_categorical_legend(
                 fontsize=legend_fontsize,
                 path_effects=legend_fontoutline,
             )
+    elif legend_loc not in {None, "none"}:
+        for label in cats:
+            ax.scatter([], [], c=palette[label], label=label)
+        if legend_loc == "right margin":
+            ax.legend(
+                frameon=False,
+                loc="center left",
+                bbox_to_anchor=(1, 0.5),
+                ncol=(1 if len(cats) <= 14 else 2 if len(cats) <= 30 else 3),
+                fontsize=legend_fontsize,
+            )
+        else:
+            ax.legend(loc=legend_loc, fontsize=legend_fontsize)
 
 
 def _get_basis(adata: AnnData, basis: str) -> np.ndarray:

src/scanpy/plotting/_docs.py

@@ -42,6 +42,23 @@
 _FontSize = Literal[
     "xx-small", "x-small", "small", "medium", "large", "x-large", "xx-large"
 ]
+_LegendLoc = Literal[
+    "none",
+    "right margin",
+    "on data",
+    "on data export",
+    "best",
+    "upper right",
+    "upper left",
+    "lower left",
+    "lower right",
+    "right",
+    "center left",
+    "center right",
+    "lower center",
+    "upper center",
+    "center",
+]
 ColorLike = Union[str, tuple[float, ...]]
 
 

src/scanpy/plotting/_tools/paga.py

@@ -2,6 +2,7 @@
 
 from functools import partial
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import matplotlib as mpl
 import matplotlib.pyplot as plt
@@ -15,6 +16,10 @@
 import scanpy as sc
 from testing.scanpy._helpers.data import pbmc3k_processed
 
+if TYPE_CHECKING:
+    from scanpy.plotting._utils import _LegendLoc
+
+
 HERE: Path = Path(__file__).parent
 ROOT = HERE / "_images"
 
@@ -111,10 +116,10 @@ def plotfunc(request):
 
 
 @pytest.fixture(
-    params=["on data", "right margin", None],
-    ids=["legend.on_data", "legend.on_right", "legend.off"],
+    params=["on data", "right margin", "lower center", None],
+    ids=["legend.on_data", "legend.on_right", "legend.on_bottom", "legend.off"],
 )
-def legend_loc(request):
+def legend_loc(request) -> _LegendLoc | None:
     return request.param
 
 
@@ -155,7 +160,7 @@ def vbounds(request):
 
 def test_missing_values_categorical(
     *,
-    fixture_request,
+    fixture_request: pytest.FixtureRequest,
     image_comparer,
     adata,
     plotfunc,
@@ -182,7 +187,13 @@ def test_missing_values_categorical(
 
 
 def test_missing_values_continuous(
-    *, fixture_request, image_comparer, adata, plotfunc, na_color, legend_loc, vbounds
+    *,
+    fixture_request: pytest.FixtureRequest,
+    image_comparer,
+    adata,
+    plotfunc,
+    na_color,
+    vbounds,
 ):
     save_and_compare_images = partial(image_comparer, MISSING_VALUES_ROOT, tol=15)
 
@@ -191,7 +202,6 @@ def test_missing_values_continuous(
     # Passing through a dict so it's easier to use default values
     kwargs = {}
     kwargs.update(vbounds)
-    kwargs["legend_loc"] = legend_loc
     if na_color is not None:
         kwargs["na_color"] = na_color