IPython line magic to generate numpy-style docstring stubs.

docstring_magic.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from IPython.core.magic import Magics, magics_class, line_magic\n",
    "import inspect, re\n",
    "\n",
    "# The class MUST call this class decorator at creation time\n",
    "@magics_class\n",
    "class TahMagics(Magics):\n",
    "    @line_magic\n",
    "    def docstring(self, line):\n",
    "        \"\"\"\n",
    "        Generate a numpy-style docstring.\n",
    "        \n",
    "        The magic supports module-level functions, classes and \n",
    "        instance methods. It generates documentation stubs for\n",
    "        arguments, attributes and exceptions.\n",
    "        \n",
    "        See https://github.com/numpy/numpy/blob/master/doc/example.py\n",
    "        for reference.\n",
    "        \"\"\"\n",
    "        # Try and get the object from the shell\n",
    "        global_ns = self.shell.user_global_ns\n",
    "        local_ns = self.shell.user_ns\n",
    "        try:\n",
    "            obj = eval(line, global_ns, local_ns)\n",
    "        except Exception as ex:\n",
    "            print \"Object `{}` not found: {}\".format(line, ex)\n",
    "            return\n",
    "        \n",
    "        # Get the function code or the __init__ function code\n",
    "        # if a class was passed\n",
    "        if inspect.isfunction(obj):\n",
    "            func_code = obj.func_code\n",
    "        elif inspect.isclass(obj):\n",
    "            func_code = obj.__init__.im_func.func_code\n",
    "        elif inspect.ismethod(obj):\n",
    "            func_code = obj.im_func.func_code\n",
    "        else:\n",
    "            raise NotImplementedError(\"Type `{}` is not implemented.\".format(type(obj)))\n",
    "            \n",
    "        # Get the arguments\n",
    "        args = inspect.getargs(func_code)\n",
    "        # Get the 'self' argument for class functions or set to None if module level\n",
    "        _self = None if inspect.isfunction(obj) else args.args.pop(0)\n",
    "\n",
    "        # Create a list of strings to concatenate\n",
    "        strs = ['Single line description.']\n",
    "            \n",
    "        # Process the parameters\n",
    "        str_args = args.args\n",
    "        if args.varargs:\n",
    "            str_args.append(\"*{}\".format(args.varargs))\n",
    "        if args.keywords:\n",
    "            str_args.append(\"**{}\".format(args.keywords))\n",
    "        if len(str_args) > 0:\n",
    "            str_args = \"\\n\".join([\"{} : {{type}}\\n    {{desc}}\".format(arg) for arg in str_args])\n",
    "            str_args = \"Parameters\\n----------\\n{}\".format(str_args)\n",
    "            strs.append(str_args)\n",
    "        \n",
    "        # Attributes\n",
    "        src = inspect.getsource(func_code)\n",
    "        if inspect.isclass(obj):\n",
    "            str_attrs = re.findall(r\"{}\\.([\\w]*)\\s*=[^=]\".format(_self), src)\n",
    "            if len(str_attrs) > 0:\n",
    "                str_attrs = \"\\n\".join([\"{} : {{type}}\\n    {{desc}}\".format(arg) for arg in str_attrs])\n",
    "                str_attrs = \"Attributes\\n----------\\n{}\".format(str_attrs)\n",
    "                strs.append(str_attrs)\n",
    "        \n",
    "        # Are there any return values?\n",
    "        if src.find(\"return\") != -1:\n",
    "            strs.append(\"Returns\\n-------\\n{{type}}\\n    {{desc}}\")\n",
    "            \n",
    "        # Are there any exceptions\n",
    "        if src.find(\"raise\") != -1 or src.find(\"assert\") != -1:\n",
    "            strs.append(\"Raises\\n------\\n{{type}}\\n    {{desc}}\")\n",
    "            \n",
    "        strs = \"\\n\\n\".join(strs)\n",
    "        print strs\n",
    "\n",
    "# Register the magic with IPython\n",
    "ip = get_ipython()\n",
    "ip.register_magics(TahMagics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Single line description.\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "arg1 : {type}\n",
      "    {desc}\n",
      "arg2 : {type}\n",
      "    {desc}\n",
      "*args : {type}\n",
      "    {desc}\n",
      "**kwargs : {type}\n",
      "    {desc}\n",
      "\n",
      "Returns\n",
      "-------\n",
      "{{type}}\n",
      "    {{desc}}\n",
      "\n",
      "Raises\n",
      "------\n",
      "{{type}}\n",
      "    {{desc}}\n"
     ]
    }
   ],
   "source": [
    "# Create a module-level function\n",
    "def module_level(arg1, arg2, *args, **kwargs):\n",
    "    if arg2 == 0:\n",
    "        raise ZeroDivisionError\n",
    "    return arg1 / arg2\n",
    "\n",
    "%docstring module_level"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Single line description.\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "arg1 : {type}\n",
      "    {desc}\n",
      "arg2 : {type}\n",
      "    {desc}\n",
      "\n",
      "Attributes\n",
      "----------\n",
      "arg1 : {type}\n",
      "    {desc}\n",
      "arg2 : {type}\n",
      "    {desc}\n"
     ]
    }
   ],
   "source": [
    "# Create a class\n",
    "class Class:\n",
    "    def __init__(self, arg1, arg2):\n",
    "        self.arg1 = arg1\n",
    "        self.arg2 = arg2\n",
    "        \n",
    "%docstring Class"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}