numba: @vectorize-decorated functions are not found by sphinx

Hi @HPLegion, thanks for the extension! Any idea on how to extend it to support @guvectorize-decorated functions?

Here’s my version of the extension that handles @guvectorized functions, if anyone’s interested:

https://github.com/legend-exp/pygama/blob/f980faa5322aac1de12b5817b1d33c3248e399b5/docs/source/extensions/numbadoc.py

At the moment we are satisfied with how it works. However, I would prefer to understand the flaky behavior, so I will probably bug the people on discourse 😛

Thank you for all your help 😃

Hi Stefano (and others)

I just wanted to give you a heads-up that I won’t be able to follow up on these issues myself for a while, unfortunately. I would like to offer more help, but time is limited at the moment.

I expect documentation to fail without the plugin (as you observed), since Sphinx gets confused by numba’s decorators. I don’t know why you see the flaky behviour in some cases though.

When I first looked into this, I really had to debug the Sphinx autodoc run to see why the documenters did not pick up on the decorated functions, and then I tried to come up with something that would be able to identify the decorated functions as more or less normally documentable objects. Then I used the chance to sprinkle in a few extras, maybe some of that broke over the last years. It is not particularly fun, but maybe a way for follow up investigations.

In the meantime I would encourage you to write on https://numba.discourse.group/ - there is a chance you will be able to get in touch with some knowledgable people who are not aware of this issue / disucussion. Judging by how lively this thread still is, this seems to affect a number of users. Maybe someone has the time and energy to really dive into this and come up with a good and maintainable lonterm solution 😃

Best of luck!

I am buried in other work at the moment, but with the collection as it is, this almost looks as if it might warrant releasing a micro-package that people can install or vendor directly into their codebase. If anyone feels the urge to create such a thing, I would certainly not hold a grudge 😛

Just had a brief look at the issue @gmarkall cross linked yesterday, seems like there are still issues linked to the missing magic wrapper information. Is there a dedicated issue for that?

@HPLegion, thanks for the sphinx extension for autodoc. I tried to change it to work with autosummary + @intrinsic:

import numba
from typing import List, Optional
from sphinx.ext.autodoc import FunctionDocumenter


class NumbaDocumenter(FunctionDocumenter):
    """Sphinx directive that also understands numba decorators"""

    objtype = 'function'  # This must be assigned to 'function' to work!!
    directivetype = 'function'

    @classmethod
    def can_document_member(cls, member, membername, isattr, parent):
        if isinstance(member, numba.core.extending._Intrinsic):
            return True
        return super(RBCDocumenter, cls).can_document_member(member, membername, isattr, parent)


def setup(app):
    """Setup Sphinx extension."""

    # Register the new documenters and directives (autojitfun, autovecfun)
    # Set the default prefix which is printed in front of the function signature
    app.setup_extension('sphinx.ext.autodoc')
    app.add_autodocumenter(NumbaDocumenter)

    return {
        'parallel_read_safe': True
    }

Also, in order for this to work, one needs to update the intrinsic wrapper:

import functools
functools.update_wrapper(my_intrinsic, my_intrinsic._defn)  # _defn is the inner function

I tried adapting the jit-documenter stuff but I couldn’t get it to work. I adopted a solution that I am ashamed of but that works: running this before sphinx 😃

find mypackage/ -name '*.py' -exec sed -i -e 's/@cuda/#@cuda/g' {} \;

I had some time yesterday and found myself thinking that bending the inner workings of numba for the sake of sphinx may not be the way to go (even though there may or may not be other advantages to this). So I had a bit of a look for ways to fixing my problem and got inspired by celery - they are offering a small sphinx extension as part of their package to allow auto-documentation of Celery (decorated) tasks.

The advantage to such an approach could be that it actually offers the possibility to customise the documentation behaviour. I hacked together a proof of concept, and it is not too terrifying, luckily Sphinx has a fairly comfortable extension API (but debugging event-driven programs turns out to be a nightmare 😨).

Here is a plaintext example of what could be done in principle. The documenter (see below) adds a prefix (@numba.vectorize) to the signature and injects information about the precompiled signature into the docstring.

@numba.vectorize ebisim.plasma.clog_ei(Ni, Ne, kbTi, kbTe, Ai, qi)

      *Precompiled signatures:[‘dddddl->d’]*

   The coulomb logarithm for ion electron collisions.

   Parameters:
      * **Ni** (*float** or **numpy.ndarray*) – <1/m^3> Ion density.

      * **Ne** (*float** or **numpy.ndarray*) – <1/m^3> Electron
        density.

      * **kbTi** (*float** or **numpy.ndarray*) – <eV> Ion
        temperature.

      * **kbTe** (*float** or **numpy.ndarray*) – <eV> Electron
        temperature.

      * **Ai** (*float** or **numpy.ndarray*) – Ion mass number.

      * **qi** (*int** or **numpy.ndarray*) – Ion charge state.

   Returns:
      Ion electron coulomb logarithm.

   Return type:
      float or numpy.ndarray

Is an extension like this something that could be interesting to feature in numba, or is it just too likely to become a maintenance liability? 😛 (The code as it is right now would likely need some bulletproofing.)

From my side this issue can be closed, since I found a solution that is working for me but if you want to keep this open to track the feature request, that’s fine.

Extension code

"""
This sphinx extension aims to improve the documentation of numba-decorated
functions. It it inspired by the design of Celery's sphinx extension
'celery.contrib.sphinx'.

Usage
-----

Add the extension to your :file:`conf.py` configuration module:

.. code-block:: python

    extensions = (...,
                  'numbadoc')

This extension adds two configuration fields, which determine the prefix
printed before jitted functions in the reference documentation.
Overwrite these values in your :file:`conf.py` to change this behaviour

.. code-block:: python

    #defaults
    numba_jit_prefix = '@numba.jit'
    numba_vectorize_prefix = '@numba.vectorize'

With the extension installed `autodoc` will automatically find
numba decorated objects and generate the correct docs.

If a vecotrized function with fixed signatures is found, these are injected
into the docstring.
"""
from typing import List, Iterator
from sphinx.domains.python import PyFunction
from sphinx.ext.autodoc import FunctionDocumenter

from numba.core.dispatcher import Dispatcher
from numba.np.ufunc.dufunc import DUFunc

class NumbaFunctionDocumenter(FunctionDocumenter):
    """Document numba decorated functions."""

    def import_object(self) -> bool:
        """Import the object given by *self.modname* and *self.objpath* and set
        it as *self.object*.

        Returns True if successful, False if an error occurred.
        """
        success = super().import_object()
        if success:
            # Store away numba wrapper
            self.jitobj = self.object
            # And bend references to underlying python function
            if hasattr(self.object, "py_func"):
                self.object = self.object.py_func
            elif hasattr(self.object, "_dispatcher") and \
                 hasattr(self.object._dispatcher, "py_func"):
                self.object = self.object._dispatcher.py_func
            else:
                success = False
        return success

    def process_doc(self, docstrings: List[List[str]]) -> Iterator[str]:
        """Let the user process the docstrings before adding them."""
        # Essentially copied from FunctionDocumenter
        for docstringlines in docstrings:
            if self.env.app:
                # let extensions preprocess docstrings
                # need to manually set 'what' to FunctionDocumenter.objtype
                # to not confuse preprocessors like napoleon with an objtype
                # that they don't know
                self.env.app.emit('autodoc-process-docstring',
                                  FunctionDocumenter.objtype,
                                  self.fullname, self.object,
                                  self.options, docstringlines)
            # This block inserts information about precompiled signatures
            # if this is a precompiled vectorized function
            if getattr(self.jitobj, "types", []) and \
               getattr(self.jitobj, "_frozen", False):
                s = "| *Precompiled signatures:" + str(self.jitobj.types) + "*"
                docstringlines.insert(0, s)
            yield from docstringlines


class JitDocumenter(NumbaFunctionDocumenter):
    """Document jit/njit decorated functions."""

    objtype = 'jitfun'

    @classmethod
    def can_document_member(cls, member, membername, isattr, parent):
        return isinstance(member, Dispatcher) and hasattr(member, 'py_func')


class VectorizeDocumenter(NumbaFunctionDocumenter):
    """Document vectorize decorated functions."""

    objtype = 'vecfun'

    @classmethod
    def can_document_member(cls, member, membername, isattr, parent):
        return isinstance(member, DUFunc) and \
               hasattr(member, '_dispatcher') and \
               hasattr(member._dispatcher, 'py_func')


class JitDirective(PyFunction):
    """Sphinx jitfun directive."""

    def get_signature_prefix(self, sig):
        return self.env.config.numba_jit_prefix


class VectorizeDirective(PyFunction):
    """Sphinx vecfun directive."""

    def get_signature_prefix(self, sig):
        return self.env.config.numba_vectorize_prefix


def setup(app):
    """Setup Sphinx extension."""
    # Register the new documenters and directives (autojitfun, autovecfun)
    # Set the default prefix which is printed in front of the function signature
    app.setup_extension('sphinx.ext.autodoc')
    app.add_autodocumenter(JitDocumenter)
    app.add_directive_to_domain('py', 'jitfun', JitDirective)
    app.add_config_value('numba_jit_prefix', '@numba.jit', True)
    app.add_autodocumenter(VectorizeDocumenter)
    app.add_directive_to_domain('py', 'vecfun', VectorizeDirective)
    app.add_config_value('numba_vectorize_prefix', '@numba.vectorize', True)

    return {
        'parallel_read_safe': True
    }

(ping @esc )