chainable scanpy methods

scanpy_chain.py

import scanpy as sc
from anndata import AnnData
from functools import partial, partialmethod, wraps
from inspect import isfunction, getmembers

class Mapper(object):
    def __init__(self, parent, mod):
        self._parent = parent
        self._mod = mod
        self._modfuncs = getmembers(mod, isfunction)
        for name, f in self._modfuncs:
            add_method(self, self._parent)(f)

def add_method(obj, firstarg=None):
    if firstarg is None:
        firstarg = obj
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            ret = partial(func, firstarg)(*args, **kwargs)
            if ret is None:  # Hacky workaround for chaining inplace operations
                ret = firstarg
            return ret
        setattr(obj, func.__name__, wrapper)
        # Not binding func, but a wrapper which acts like it
        return func # returning func means func can still be used normally, only meaningful for decorator usecase
    return decorator

AnnData.pl = property(partial(Mapper, mod=sc.pl))
AnnData.tl = property(partial(Mapper, mod=sc.tl))
AnnData.pp = property(partial(Mapper, mod=sc.pp))

# Example

pbmc = sc.datasets.pbmc3k()
pbmc = (pbmc
    .pp.normalize_per_cell(counts_per_cell_after=1000)
    .pp.log1p()
    .pp.pca()
)
pbmc.pl.pca()


#######################
# Extension: subsetting
#######################

def select(adata, obs=None, var=None):
    is_var = var is None
    is_obs = obs is None
    if is_var and is_obs:
        raise ValueError()
    if obs:
        idx = adata.obs.loc[obs].index
        return adata[idx, :]
    if var:
        idx = adata.var.loc[var].index
        return adata[:, idx]

# Basically the same:
# setattr(AnnData, "select", select)
add_method(AnnData)(select)

# Example

pbmc = sc.datasets.pbmc68k_reduced()
(pbmc
    .select(obs=lambda x: x["bulk_labels"] == "Dendritic")
    .pl.umap()
)

Would be improved by having scanpy methods return the modified object if performed inplace, and by allowing multiple levels of mapper. Probably would be good to whitelist which methods are available through mapper. Alternatively, could check which ones take an AnnData for the first arg.