Search code examples
pythoncoding-styledocstringfirst-class-functionsnumpy-ufunc

Documenting first-class-assigned functions


I have a function defined by making use of the first-class nature of Python functions, as follows:

add_relative = np.frompyfunc(lambda a, b: (1 + a) * (1 + b) - 1, 2, 1)

Either I need a way to add a docstring to the function defined as it is, or achieve the same thing using the more common format, so that I can write a docstring in the normal way:

def add_relative(a, b):
    """
    Docstring
    """
    return np.frompyfunc(lambda a, b: (1 + a) * (1 + b) - 1, 2, 1)(a, b)

which works when the function is called like

 add_relative(arr1, arr2)

but I then lose the ability to call methods, for example

add_relative.accumulate(foo_arr, dtype=np.object)

I guess this is because the function becomes more like a class when using frompyfunc, being derived from ufunc.

I'm thinking I might need to define a class, rather than a function, but I'm not sure how. I would be ok with that because then I can easily add a docstring as normal.

I tagged this coding-style because the original method works but simply can't be easily documented, and I'm sorry if the title is not clear, I don't know the correct vocabulary to describe this.


Solution

  • Update 1: Close, but this still isn't good enough. Because the __doc__ attribute of the decorated function cannot be updated, and because Sphinx still only picks up the docstring of the decorated function, this doesn't solve my problem.

    Update 2: The solution that I proposed below is nice for documentation within the source code. For documentation with Sphinx I ended up just overwriting the docstring with

    .. function:: sum_relative(a, b)
    
       <Docstring written in rst format>
    

    It's ugly, it's hacky and it's manual but it means that I have my nice documentation in the source code, and I have my nice documentation in Sphinx.

    All of the issues stem from the fact that the __doc__ attribute of a numpy.ufunc is immutable. If anyone knows why, I'd love to hear why. I'm guessing something related to the fact that it comes from something written in C, not pure Python. Regardless, it's very annoying.


    I found that I could solve the issue using a decorator to apply np.frompyfunc().

    I write the base function (the lambda in the original example) and add a docstring as normal and then apply the decorator:

    def as_ufunc(func):
        return np.frompyfunc(func, 2, 1)
    
    @as_ufunc
    def sum_relative(a, b):
        """
        Docstring
        """
        return (1 + a) * (1 + b) - 1
    

    It's not a perfect solution for the following reasons:

    • sum_relative.__doc__ is overwritten by frompyfunc to a generic and unhelpful docstring. I don't mind here because I really care about the docs generated with Sphinx from the docstring, and not accessing it programatically. You might think to try something like functools.wraps or functools.update_wrapper, however the __doc__ member of a numpy.ufunc is apparenly immutable.

    • I have to hardcode the second two arguments for frompyfunc. I don't mind here because all the cases that I use it here will require the same values.

    • Edit: It is possible to get around the above point, it's a little more verbose, but not much:

      def as_ufunc(nin, nout):
          def _decorator(func):
              return np.frompyfunc(func, nin, nout)
          return _decorator
      
      @as_ufunc(2, 1)
      def sum_relative(a, b):
          """
          Docstring
          """
          return (1 + a) * (1 + b) - 1
      
    • It's more verbose than the original solution. I don't mind because I now have the docstring.