Automatically add type annotations to all API functions

scanpy/api/__init__.py

@@ -28,3 +28,8 @@
 # some stuff that is not actually documented...
 from .. import utils
 from .. import rtools
+
+
+import sys
+utils.annotate_doc_types(sys.modules[__name__], 'scanpy')
+del sys

scanpy/neighbors/__init__.py

@@ -1,8 +1,12 @@
+from typing import Union, Optional, Any, Mapping, Callable
+
 import numpy as np
 import scipy
-from scipy.sparse import issparse
-from scipy.sparse import coo_matrix
+from anndata import AnnData
+from numpy.random import RandomState
+from scipy.sparse import issparse, coo_matrix
 from sklearn.metrics import pairwise_distances
+
 from .. import settings
 from .. import logging as logg
 from .. utils import doc_params
@@ -15,16 +19,17 @@
 
 @doc_params(n_pcs=doc_n_pcs, use_rep=doc_use_rep)
 def neighbors(
-        adata,
-        n_neighbors=15,
-        n_pcs=None,
-        use_rep=None,
-        knn=True,
-        random_state=0,
-        method='umap',
-        metric='euclidean',
-        metric_kwds={},
-        copy=False):
+    adata: AnnData,
+    n_neighbors: int = 15,
+    n_pcs: Optional[int] = None,
+    use_rep: Optional[str] = None,
+    knn: bool = True,
+    random_state: Optional[Union[int, RandomState]] = 0,
+    method: str = 'umap',
+    metric: Union[str, Callable[[np.ndarray, np.ndarray], float]] = 'euclidean',
+    metric_kwds: Mapping[str, Any] = {},
+    copy: bool = False
+) -> Optional[AnnData]:
     """\
     Compute a neighborhood graph of observations [McInnes18]_.
 
@@ -36,9 +41,9 @@ def neighbors(
 
     Parameters
     ----------
-    adata : :class:`~anndata.AnnData`
+    adata
         Annotated data matrix.
-    n_neighbors : `int`, optional (default: 15)
+    n_neighbors
         The size of local neighborhood (in terms of number of neighboring data
         points) used for manifold approximation. Larger values result in more
         global views of the manifold, while smaller values result in more local
@@ -48,15 +53,21 @@ def neighbors(
         `n_neighbors` neighbor.
     {n_pcs}
     {use_rep}
-    knn : `bool`, optional (default: `True`)
+    knn
         If `True`, use a hard threshold to restrict the number of neighbors to
         `n_neighbors`, that is, consider a knn graph. Otherwise, use a Gaussian
         Kernel to assign low weights to neighbors more distant than the
         `n_neighbors` nearest neighbor.
+    random_state
+        A numpy random seed.
     method : {{'umap', 'gauss', `None`}}  (default: `'umap'`)
         Use 'umap' [McInnes18]_ or 'gauss' (Gauss kernel following [Coifman05]_
         with adaptive width [Haghverdi16]_) for computing connectivities.
-    copy : `bool` (default: `False`)
+    metric
+        A known metric’s name or a callable that returns a distance.
+    metric_kwds
+        Options for the metric.
+    copy
         Return a copy instead of writing to adata.
 
     Returns
@@ -459,7 +470,7 @@ def _backwards_compat_get_full_eval(adata):
         return adata.uns['diffmap_evals']
 
 
-class OnFlySymMatrix():
+class OnFlySymMatrix:
     """Emulate a matrix where elements are calculated on the fly.
     """
     def __init__(self, get_row, shape, DC_start=0, DC_end=-1, rows=None, restrict_array=None):
@@ -503,19 +514,21 @@ def restrict(self, index_array):
                               rows=self.rows, restrict_array=index_array)
 
 
-class Neighbors():
+class Neighbors:
     """Data represented as graph of nearest neighbors.
 
     Represent a data matrix as a graph of nearest neighbor relations (edges)
     among data points (nodes).
 
     Parameters
     ----------
-    adata : :class:`~anndata.AnnData`
-        An annotated data matrix.
+    adata
+        Annotated data object.
+    n_dcs
+        Number of diffusion components to use.
     """
 
-    def __init__(self, adata, n_dcs=None):
+    def __init__(self, adata: AnnData, n_dcs: Optional[int] = None):
         self._adata = adata
         self._init_iroot()
         # use the graph in adata
@@ -646,25 +659,25 @@ def to_igraph(self):
 
     @doc_params(n_pcs=doc_n_pcs, use_rep=doc_use_rep)
     def compute_neighbors(
-            self,
-            n_neighbors=30,
-            knn=True,
-            n_pcs=None,
-            use_rep=None,
-            method='umap',
-            random_state=0,
-            write_knn_indices=False,
-            precompute_metric=None,
-            metric='euclidean',
-            metric_kwds={}):
+        self,
+        n_neighbors: int = 30,
+        knn: bool = True,
+        n_pcs: Optional[int] = None,
+        use_rep: Optional[str] = None,
+        method: str = 'umap',
+        random_state: Optional[Union[RandomState, int]] = 0,
+        write_knn_indices: bool = False,
+        metric: str = 'euclidean',
+        metric_kwds: Mapping[str, Any] = {}
+    ) -> None:
         """\
         Compute distances and connectivities of neighbors.
 
         Parameters
         ----------
-        n_neighbors : `int`, optional (default: 30)
+        n_neighbors
              Use this number of nearest neighbors.
-        knn : `bool`, optional (default: `True`)
+        knn
              Restrict result to `n_neighbors` nearest neighbors.
         {n_pcs}
         {use_rep}

scanpy/utils.py

@@ -1,18 +1,64 @@
 """Utility functions and classes
 """
 
+import inspect
 from collections import namedtuple
+from functools import partial
+from types import ModuleType, MethodType, FunctionType
+from typing import Iterable, Union, Callable
+
 import numpy as np
 import scipy.sparse
 from natsort import natsorted
 from textwrap import dedent
 from pandas.api.types import CategoricalDtype
+try:
+    from numpydoc import docscrape_sphinx
+except ImportError:
+    docscrape_sphinx = None
 
 from . import settings
 from . import logging as logg
 
 EPS = 1e-15
 
+def getdoc(c_or_f: Union[Callable, type]):
+    doc = docscrape_sphinx.get_doc_object(c_or_f)
+    if isinstance(doc, docscrape_sphinx.ClassDoc) and hasattr(c_or_f, '__init__'):
+        sig = inspect.signature(c_or_f.__init__)
+    else:
+        sig = inspect.signature(c_or_f)
+    if isinstance(doc, docscrape_sphinx.FunctionDoc):
+        def type_doc(name):
+            param = sig.parameters[name]  # type: inspect.Parameter
+            cls = getattr(param.annotation, '__qualname__', repr(param.annotation))
+            return '{}, optional (default: {!r})'.format(cls, param.default) if param.default else cls
+        doc['Parameters'] = [
+            (name, typ if typ else type_doc(name), docs)
+            for name, typ, docs in doc['Parameters']
+        ]
+    return str(doc)
+
+
+def descend_classes_and_funcs(mod: ModuleType, root: str):
+    for obj in vars(mod).values():
+        if not getattr(obj, '__module__', getattr(obj, '__qualname__', getattr(obj, '__name__', ''))).startswith(root):
+            continue
+        if isinstance(obj, Callable):
+            yield obj
+            if isinstance(obj, type):
+                yield from (m for m in vars(obj).values() if isinstance(m, Callable))
+        elif isinstance(obj, ModuleType):
+            yield from descend_classes_and_funcs(obj, root)
+
+
+def annotate_doc_types(mod: ModuleType, root: str):
+    if docscrape_sphinx is None:
+        return
+    for c_or_f in descend_classes_and_funcs(mod, root):
+        c_or_f.getdoc = partial(getdoc, c_or_f)
+
+
 def doc_params(**kwds):
     """\
     Docstrings should start with "\" in the first line for proper formatting.