# Buggy code
def _plot_categories_as_colorblocks(
    groupby_ax: Axes,
    obs_tidy: pd.DataFrame,
    colors=None,
    orientation: Literal["top", "bottom", "left", "right"] = "left",
    cmap_name: str = "tab20",
):
    """\
    Plots categories as colored blocks. If orientation is 'left', the categories
    are plotted vertically, otherwise they are plotted horizontally.

    Parameters
    ----------
    groupby_ax
    obs_tidy
    colors
        Sequence of valid color names to use for each category.
    orientation
    cmap_name
        Name of colormap to use, in case colors is None

    Returns
    -------
    ticks position, labels, colormap
    """

    groupby = obs_tidy.index.name
    from matplotlib.colors import ListedColormap, BoundaryNorm

    if colors is None:
        groupby_cmap = pl.get_cmap(cmap_name)
    else:
        groupby_cmap = ListedColormap(colors, groupby + "_cmap")
    norm = BoundaryNorm(np.arange(groupby_cmap.N + 1) - 0.5, groupby_cmap.N)

    # determine groupby label positions such that they appear
    # centered next/below to the color code rectangle assigned to the category
    value_sum = 0
    ticks = []  # list of centered position of the labels
    labels = []
    label2code = {}  # dictionary of numerical values asigned to each label
    for code, (label, value) in enumerate(
        obs_tidy.index.value_counts(sort=False).iteritems()
    ):
        ticks.append(value_sum + (value / 2))
        labels.append(label)
        value_sum += value
        label2code[label] = code

    groupby_ax.grid(False)

    if orientation == "left":
        groupby_ax.imshow(
            np.array([[label2code[lab] for lab in obs_tidy.index]]).T,
            aspect="auto",
            cmap=groupby_cmap,
            norm=norm,
        )
        if len(labels) > 1:
            groupby_ax.set_yticks(ticks)
            groupby_ax.set_yticklabels(labels)

        # remove y ticks
        groupby_ax.tick_params(axis="y", left=False, labelsize="small")
        # remove x ticks and labels
        groupby_ax.tick_params(axis="x", bottom=False, labelbottom=False)

        # remove surrounding lines
        groupby_ax.spines["right"].set_visible(False)
        groupby_ax.spines["top"].set_visible(False)
        groupby_ax.spines["left"].set_visible(False)
        groupby_ax.spines["bottom"].set_visible(False)

        groupby_ax.set_ylabel(groupby)
    else:
        groupby_ax.imshow(
            np.array([[label2code[lab] for lab in obs_tidy.index]]),
            aspect="auto",
            cmap=groupby_cmap,
            norm=norm,
        )
        if len(labels) > 1:
            groupby_ax.set_xticks(ticks)
            if max([len(x) for x in labels]) < 3:
                # if the labels are small do not rotate them
                rotation = 0
            else:
                rotation = 90
            groupby_ax.set_xticklabels(labels, rotation=rotation)

        # remove x ticks
        groupby_ax.tick_params(axis="x", bottom=False, labelsize="small")
        # remove y ticks and labels
        groupby_ax.tick_params(axis="y", left=False, labelleft=False)

        # remove surrounding lines
        groupby_ax.spines["right"].set_visible(False)
        groupby_ax.spines["top"].set_visible(False)
        groupby_ax.spines["left"].set_visible(False)
        groupby_ax.spines["bottom"].set_visible(False)

        groupby_ax.set_xlabel(groupby)

    return label2code, ticks, labels, groupby_cmap, norm


# Traceback (TypeError)
TypeError                                 Traceback (most recent call last)
<ipython-input-56-f1ba710dac43> in <module>
9
10 sc.pl.heatmap(a, var_names=a.var_names[:5], groupby=None, swap_axes=True)  # is shifted to the left (fig. 1)
---> 11 sc.pl.heatmap(a, var_names=a.var_names[:5], groupby="foo", swap_axes=True)
12
13 a.X[:16, :] = 0

~/.miniconda3/envs/cellrank/lib/python3.8/site-packages/scanpy/plotting/_anndata.py in heatmap(adata, var_names, groupby, use_raw, log, num_categories, dendrogram, gene_symbols, var_group_positions, var_group_labels, var_group_rotation, layer, standard_scale, swap_axes, show_gene_labels, show, save, figsize, **kwds)
1220                 groupby_cmap,
1221                 norm,
-> 1222             ) = _plot_categories_as_colorblocks(
1223                 groupby_ax, obs_tidy, colors=groupby_colors, orientation='bottom'
1224             )

~/.miniconda3/envs/cellrank/lib/python3.8/site-packages/scanpy/plotting/_anndata.py in _plot_categories_as_colorblocks(groupby_ax, obs_tidy, colors, orientation, cmap_name)
2361         if len(labels) > 1:
2362             groupby_ax.set_xticks(ticks)
-> 2363             if max([len(x) for x in labels]) < 3:
2364                 # if the labels are small do not rotate them
2365                 rotation = 0

~/.miniconda3/envs/cellrank/lib/python3.8/site-packages/scanpy/plotting/_anndata.py in <listcomp>(.0)
2361         if len(labels) > 1:
2362             groupby_ax.set_xticks(ticks)
-> 2363             if max([len(x) for x in labels]) < 3:
2364                 # if the labels are small do not rotate them
2365                 rotation = 0

TypeError: object of type 'int' has no len()

# Fixed code
def _plot_categories_as_colorblocks(
    groupby_ax: Axes,
    obs_tidy: pd.DataFrame,
    colors=None,
    orientation: Literal["top", "bottom", "left", "right"] = "left",
    cmap_name: str = "tab20",
):
    """\
    Plots categories as colored blocks. If orientation is 'left', the categories
    are plotted vertically, otherwise they are plotted horizontally.

    Parameters
    ----------
    groupby_ax
    obs_tidy
    colors
        Sequence of valid color names to use for each category.
    orientation
    cmap_name
        Name of colormap to use, in case colors is None

    Returns
    -------
    ticks position, labels, colormap
    """

    groupby = obs_tidy.index.name
    from matplotlib.colors import ListedColormap, BoundaryNorm

    if colors is None:
        groupby_cmap = pl.get_cmap(cmap_name)
    else:
        groupby_cmap = ListedColormap(colors, groupby + "_cmap")
    norm = BoundaryNorm(np.arange(groupby_cmap.N + 1) - 0.5, groupby_cmap.N)

    # determine groupby label positions such that they appear
    # centered next/below to the color code rectangle assigned to the category
    value_sum = 0
    ticks = []  # list of centered position of the labels
    labels = []
    label2code = {}  # dictionary of numerical values asigned to each label
    for code, (label, value) in enumerate(
        obs_tidy.index.value_counts(sort=False).iteritems()
    ):
        ticks.append(value_sum + (value / 2))
        labels.append(label)
        value_sum += value
        label2code[label] = code

    groupby_ax.grid(False)

    if orientation == "left":
        groupby_ax.imshow(
            np.array([[label2code[lab] for lab in obs_tidy.index]]).T,
            aspect="auto",
            cmap=groupby_cmap,
            norm=norm,
        )
        if len(labels) > 1:
            groupby_ax.set_yticks(ticks)
            groupby_ax.set_yticklabels(labels)

        # remove y ticks
        groupby_ax.tick_params(axis="y", left=False, labelsize="small")
        # remove x ticks and labels
        groupby_ax.tick_params(axis="x", bottom=False, labelbottom=False)

        # remove surrounding lines
        groupby_ax.spines["right"].set_visible(False)
        groupby_ax.spines["top"].set_visible(False)
        groupby_ax.spines["left"].set_visible(False)
        groupby_ax.spines["bottom"].set_visible(False)

        groupby_ax.set_ylabel(groupby)
    else:
        groupby_ax.imshow(
            np.array([[label2code[lab] for lab in obs_tidy.index]]),
            aspect="auto",
            cmap=groupby_cmap,
            norm=norm,
        )
        if len(labels) > 1:
            groupby_ax.set_xticks(ticks)
            if max([len(str(x)) for x in labels]) < 3:
                # if the labels are small do not rotate them
                rotation = 0
            else:
                rotation = 90
            groupby_ax.set_xticklabels(labels, rotation=rotation)

        # remove x ticks
        groupby_ax.tick_params(axis="x", bottom=False, labelsize="small")
        # remove y ticks and labels
        groupby_ax.tick_params(axis="y", left=False, labelleft=False)

        # remove surrounding lines
        groupby_ax.spines["right"].set_visible(False)
        groupby_ax.spines["top"].set_visible(False)
        groupby_ax.spines["left"].set_visible(False)
        groupby_ax.spines["bottom"].set_visible(False)

        groupby_ax.set_xlabel(groupby)

    return label2code, ticks, labels, groupby_cmap, norm