IPython line magic to generate numpy-style docstring stubs.

docstring_magic.py

from IPython.core.magic import Magics, magics_class, line_magic
import inspect, re

# The class MUST call this class decorator at creation time
@magics_class
class TahMagics(Magics):
    @line_magic
    def docstring(self, line):
        """
        Generate a numpy-style docstring.
        
        The magic supports module-level functions, classes and 
        instance methods. It generates documentation stubs for
        arguments, attributes and exceptions.
        
        See https://github.com/numpy/numpy/blob/master/doc/example.py
        for reference.
        """
        # Try and get the object from the shell
        global_ns = self.shell.user_global_ns
        local_ns = self.shell.user_ns
        try:
            obj = eval(line, global_ns, local_ns)
        except Exception as ex:
            print "Object `{}` not found: {}".format(line, ex)
            return
        
        # Get the function code or the __init__ function code
        # if a class was passed
        if inspect.isfunction(obj):
            func_code = obj.func_code
        elif inspect.isclass(obj):
            func_code = obj.__init__.im_func.func_code
        elif inspect.ismethod(obj):
            func_code = obj.im_func.func_code
        else:
            raise NotImplementedError("Type `{}` is not implemented.".format(type(obj)))
            
        # Get the arguments
        args = inspect.getargs(func_code)
        # Get the 'self' argument for class functions or set to None if module level
        _self = None if inspect.isfunction(obj) else args.args.pop(0)

        # Create a list of strings to concatenate
        strs = ['Single line description.']
            
        # Process the parameters
        str_args = args.args
        if args.varargs:
            str_args.append("*{}".format(args.varargs))
        if args.keywords:
            str_args.append("**{}".format(args.keywords))
        if len(str_args) > 0:
            str_args = "\n".join(["{} : {{type}}\n    {{desc}}".format(arg) for arg in str_args])
            str_args = "Parameters\n----------\n{}".format(str_args)
            strs.append(str_args)
        
        # Attributes
        src = inspect.getsource(func_code)
        if inspect.isclass(obj):
            str_attrs = re.findall(r"{}\.([\w]*)\s*=[^=]".format(_self), src)
            if len(str_attrs) > 0:
                str_attrs = "\n".join(["{} : {{type}}\n    {{desc}}".format(arg) for arg in str_attrs])
                str_attrs = "Attributes\n----------\n{}".format(str_attrs)
                strs.append(str_attrs)
        
        # Are there any return values?
        if src.find("return") != -1:
            strs.append("Returns\n-------\n{{type}}\n    {{desc}}")
            
        # Are there any exceptions
        if src.find("raise") != -1 or src.find("assert") != -1:
            strs.append("Raises\n------\n{{type}}\n    {{desc}}")
            
        strs = "\n\n".join(strs)
        print strs

# Register the magic with IPython
ip = get_ipython()
ip.register_magics(TahMagics)