05_literate_programming_in_notebooks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Literate Development in Jupyter\n",
    "\n",
    "The are three goals in this notebook. First, we want to be able to import a notebook like we import modules in a python package. Second, we want to convert those notebooks into a python package. Third, we want to be able to omit some cells from the resulting python files on conversion, but still be able to execute them from the notebook. \n",
    "\n",
    "## Jupyter Notebook Importer\n",
    "\n",
    "Regular python libraries are typically kept in a 'lib' directory. The location of such a directory is stored in the systems path (accessed via `sys.path` in the `sys` module in Python)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/Users/nicolebrewer/Repos/nicole-brewer-github/literate-python',\n",
       " '/Users/nicolebrewer/miniconda3/envs/jupyter/lib/python39.zip',\n",
       " '/Users/nicolebrewer/miniconda3/envs/jupyter/lib/python3.9',\n",
       " '/Users/nicolebrewer/miniconda3/envs/jupyter/lib/python3.9/lib-dynload',\n",
       " '',\n",
       " '/Users/nicolebrewer/miniconda3/envs/jupyter/lib/python3.9/site-packages',\n",
       " '/Users/nicolebrewer/miniconda3/envs/jupyter/lib/python3.9/site-packages/IPython/extensions',\n",
       " '/Users/nicolebrewer/.ipython']"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import sys\n",
    "sys.path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But what about an imports from modules where it doesn't make sense to store a location entry on `sys.path`? [PEP 302](https://www.python.org/dev/peps/pep-0302/) provides an alternative that allows user to extend the [importlib](https://docs.python.org/3/library/importlib.html) package, which implements the Python `import` statement. And are quite a few real-life reasons why you might want to write your own.\n",
    "\n",
    "> Extending the import mechanism is needed when you want to load modules that are stored in a non-standard way. Examples include modules that are bundled together in an archive; byte code that is not stored in a pyc formatted file; modules that are loaded from a database over a network.\n",
    "\n",
    "In our case, we want to import code from our Jupyter notebooks just like we would import from a regular Python module. So how do we achieve this?\n",
    "\n",
    "PEP 302 introduced a list of import objects, `sys.meta_path` that is traversed before modules in `sys.path`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[_frozen_importlib.BuiltinImporter,\n",
       " _frozen_importlib.FrozenImporter,\n",
       " _frozen_importlib_external.PathFinder,\n",
       " <six._SixMetaPathImporter at 0x110326fa0>]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sys.meta_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The `importlib` Package\n",
    "\n",
    "The `importlib` package provides an implementation of the `import` statement, and by extension the `__import__` function, written in Python source code. It has several purposes, but in our case it is useful because it provides a way for us to write our own custom \"importer\" in Python. We will use this package to implement what is called \"importer protocol\" in order to import Jupyter Notebooks as modules.\n",
    "\n",
    "\n",
    "### The \"Importer Protocol\"\n",
    "\n",
    "The Importer Protocol operates at the level of individual imports. That is to say that all dotted imports are broken down into their components before they are imported.\n",
    "For example, by the time an importer gets a request for \"spam.ham\", module \"spam\" has already been imported. \n",
    "\n",
    "The protocal specifies two objects for importing modules: a **finder** and a **loader**. \n",
    "\n",
    "\n",
    "### Implementation\n",
    "\n",
    "The API for these objects has changed since it's original introduction, but protocol remains the same. We will use provided abstract classes to implement the most up-to-date implementation of the protocol.\n",
    "\n",
    "```\n",
    "object\n",
    " +-- Finder (deprecated)\n",
    " |    +-- MetaPathFinder\n",
    " |    +-- PathEntryFinder\n",
    " +-- Loader\n",
    "      +-- ResourceLoader --------+\n",
    "      +-- InspectLoader          |\n",
    "           +-- ExecutionLoader --+\n",
    "                                 +-- FileLoader\n",
    "                                 +-- SourceLoader\n",
    "```\n",
    "\n",
    "As stated earlier, we will be storing our import object in `sys.meta_path`. This object is called a finder. In our case we want to implement a `MetaPathFinder` class, which is [defined](https://docs.python.org/3/glossary.html#term-meta-path-finder) as a finder returned by a search of sys.meta_path.\n",
    "\n",
    "Additionally, we will implement an `ExecutionLoader` which inspects and executes Python code in order to load the module."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fullname argument of both methods is the fully qualified module name, for example \"spam.eggs.ham\". When finder.find_module(\"spam.eggs.ham\") is called, \"spam.eggs\" has already been imported and added to sys.modules, so you don't have to worry about analyzing dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib\n",
    "from IPython import get_ipython\n",
    "import nbformat\n",
    "import os, sys\n",
    "from IPython.core.interactiveshell import InteractiveShell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NotebookLoader(importlib.abc.InspectLoader, importlib.abc.MetaPathFinder):\n",
    "    '''An abstract base class for a loader which implements the optional\n",
    "    PEP 302 protocol for loaders that inspect modules.\n",
    "    \n",
    "    Parameters:\n",
    "        path: path to the directory containing the module [string]\n",
    "        module_name: a name to use for the module. Path basename is default [string]\n",
    "        remove_magics: remove all cell and line magics from Notebook cell source code\n",
    "    '''\n",
    "    def __init__(self, module_path, remove_magics=True):\n",
    "        self.remove_magics = remove_magics\n",
    "        self.shell = InteractiveShell.instance()\n",
    "        self.module_path = os.path.expanduser(module_path)  \n",
    "        if not os.path.isdir(self.module_path):\n",
    "            raise FileNotFoundError('Module directory %s not found.' % self.module_path)\n",
    "        self.module_name = [dirname for dirname in self.module_path.split('/') if dirname != ''][-1]\n",
    "\n",
    "    def get_filename(self, fullname):\n",
    "        '''An abstract method that is to return the value of __file__ for the specified module.\n",
    "        If no path is available, ImportError is raised. If source code is available, then the \n",
    "        method should return the path to the source file\n",
    "        '''\n",
    "        parts = fullname.split('.')\n",
    "        if parts[0] != self.module_name:\n",
    "            raise ImportError('Module %s doesn\\'t exist' % parts[0])\n",
    "        if len(parts) == 1:\n",
    "            nb_path = 'this file does not exist'\n",
    "            init_path = '__init__.py'\n",
    "        elif len(parts) == 2:\n",
    "            nb_path = parts[1] + '.ipynb'\n",
    "            init_path = os.path.join(parts[1], '__init__.py')\n",
    "        else:\n",
    "            parts[-1] = parts[-1] + '.ipynb'\n",
    "            nb_path = os.path.join(parts[1], *parts[2:])\n",
    "            init_path = os.path.join(parts[1], *parts[2:], '__init__.py')\n",
    "            \n",
    "        for path in (nb_path, init_path):\n",
    "            fullpath = os.path.join(self.module_path, path)\n",
    "            if os.path.exists(fullpath):\n",
    "                return fullpath\n",
    "        raise ImportError('Module %s does not exist at %s' % (fullname, fullpath))\n",
    "            \n",
    "        \n",
    "    def get_code(self, fullname):\n",
    "        '''Return the code object for a module, or None if the module does not have a code\n",
    "        object (as would be the case, for example, for a built-in module). Raise an \n",
    "        ImportError if loader cannot find the requested module.\n",
    "        '''\n",
    "        filename = self.get_filename(fullname)\n",
    "        return self.get_code_from_file(filename)\n",
    "\n",
    "    def get_source(self, fullname):\n",
    "        '''An abstract method to return the source of a module. It is returned as a\n",
    "        text string using universal newlines, translating all recognized line separators\n",
    "        into '\\n' characters. Returns None if no source is available (e.g. a built-in module).\n",
    "        Raises ImportError if the loader cannot find the module specified.\n",
    "        '''\n",
    "        return self.get_code(fullname)\n",
    "       \n",
    "    def get_code_from_file(self, filename):\n",
    "        if self.is_nb(filename):\n",
    "            return self.get_nb_source(filename)\n",
    "        else:\n",
    "            return self.get_init_source(filename)\n",
    "        \n",
    "            \n",
    "    def get_init_source(self, filename):\n",
    "        with open(filename, 'r') as f:\n",
    "            return f.read()\n",
    "        raise ModuleNotFoundError('Filename %s not readable' % filename)\n",
    "        \n",
    "    def get_nb_source(self, filename):\n",
    "        try:\n",
    "            nb = nbformat.read(filename, as_version=4)\n",
    "            code_cells = [c for c in nb['cells'] if c['cell_type'] == 'code']\n",
    "            nb_source = ''\n",
    "            for code_cell in code_cells:\n",
    "                if self.remove_magics:\n",
    "                    nb_source += self.remove_all_magics(code_cell.source) + '\\n'\n",
    "                else: # turn magics into executable code\n",
    "                    nb_source += self.shell.input_transformer_manager.transform_cell(code_cell.source)\n",
    "            return nb_source\n",
    "        except:\n",
    "            raise ModuleNotFoundExpection('Notebook %s not readable' % filename)\n",
    "\n",
    "    def remove_all_magics(self, string):\n",
    "        ret = ''\n",
    "        if string.startswith('%%'):\n",
    "            return ret\n",
    "        for line in string.splitlines():\n",
    "            if not line.lstrip().startswith('%'):\n",
    "                ret += line + '\\n'\n",
    "        return ret\n",
    "    \n",
    "    def is_nb(self, filename):\n",
    "        return filename.endswith('.ipynb')\n",
    "\n",
    "    def is_package(self, fullname):\n",
    "        '''An abstract method to return a true value if the module is a package, a false\n",
    "        value otherwise. ImportError is raised if the loader cannot find the module.\n",
    "        '''\n",
    "        return fullname == self.module_name\n",
    "    \n",
    "    def exec_module(self, module):\n",
    "        '''An abstract method that executes the module in its own namespace when a module \n",
    "        is imported or reloaded. The module should already be initialized when exec_module() \n",
    "        is called. When this method exists, create_module() must be defined.\n",
    "        '''\n",
    "        try:\n",
    "            exec(self.get_code_from_file(module.__file__), module.__dict__)\n",
    "        except:\n",
    "            print('Error executing %s' % module.__file__)\n",
    "            print(self.get_code_from_file(module.__file__))\n",
    "            raise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NotebookFinder(importlib.abc.MetaPathFinder):\n",
    "    \n",
    "    def __init__(self, module_path):\n",
    "        self.module_path = module_path\n",
    "\n",
    "    def find_spec(self, fullname, path, target=None):\n",
    "        '''An abstract method for finding a spec for the specified module. If this\n",
    "        is a top-level import, path will be None. Otherwise, this is a search for\n",
    "        a subpackage or module and path will be the value of __path__ from the parent\n",
    "        package. If a spec cannot be found, None is returned. When passed in, target\n",
    "        is a module object that the finder may use to make a more educated guess about\n",
    "        what spec to return. importlib.util.spec_from_loader() may be useful for \n",
    "        implementing concrete MetaPathFinders.\n",
    "        '''\n",
    "        loader = NotebookLoader(self.module_path) # __loader__: The Loader that should be used when loading the module. Finders should always set this.\n",
    "        name = fullname # __name__: A string for the fully-qualified name of the module.\n",
    "        # __cached__: String for where the compiled module should be stored (or None).\n",
    "        # __package__: (Read-only) The fully-qualified name of the package under which the module should be loaded as a submodule (or the empty string for top-level modules). For packages, it is the same as __name__.\n",
    "        location = loader.get_filename(fullname) # __file__: The path to where the module data is stored (not set for built-in modules).\n",
    "        self.submodule_search_locations = None # __path__: List of strings for where to find submodules, if a package (None otherwise).\n",
    "        if not location.endswith('.ipynb'):\n",
    "            (spec_dir, _) = os.path.split(location)\n",
    "            self.submodule_search_locations = []\n",
    "            self.submodule_search(spec_dir)\n",
    "        spec = importlib.util.spec_from_file_location(name, location, loader=loader, submodule_search_locations=self.submodule_search_locations)\n",
    "        return spec\n",
    "    \n",
    "    def submodule_search(self, location):\n",
    "        locations = [os.path.join(location, f) for f in os.listdir(location) if not f.startswith('.')]\n",
    "        locations = [d for d in locations if os.path.isdir(d) and os.path.isfile(os.path.join(d, '__init__.py'))]\n",
    "        self.submodule_search_locations = locations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Register the hook\n",
    "\n",
    "Now we register the NotebookFinder with sys.meta_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.meta_path.append(NotebookFinder('nbs'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<class '_frozen_importlib.BuiltinImporter'>, <class '_frozen_importlib.FrozenImporter'>, <class '_frozen_importlib_external.PathFinder'>, <six._SixMetaPathImporter object at 0x110326fa0>, <__main__.NotebookFinder object at 0x110fa4070>]\n"
     ]
    }
   ],
   "source": [
    "print(sys.meta_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can import our notebooks just like we would import a module from a package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "another_module.ipynb  child/  module.ipynb\r\n"
     ]
    }
   ],
   "source": [
    "%ls nbs/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nbs.module import hello_world"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello, world!\n"
     ]
    }
   ],
   "source": [
    "hello_world()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This way of importing is really robust, and we can make relative imports in submodules just like you would expect to be able to in a python package.\n",
    "\n",
    "For example, `nbs/another_module.ipynb` makes a relative import from `module` and calls `hello_world()`\n",
    "\n",
    "```python\n",
    "from .module import hello_world\n",
    "\n",
    "def conversation():\n",
    "    hello_world()\n",
    "    print(\"And hi to you too\")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello, world!\n",
      "And hi to you too!\n"
     ]
    }
   ],
   "source": [
    "from nbs.another_module import conversation\n",
    "\n",
    "conversation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "baby_module.ipynb\r\n"
     ]
    }
   ],
   "source": [
    "%ls nbs/child/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hello, world!\n",
      "Goo goo, gah gah\n"
     ]
    }
   ],
   "source": [
    "from nbs.child.baby_module import babble\n",
    "\n",
    "babble()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So now we have an import system that we can use with Jupyter Notebooks. Amazing! This is a great step toward a literate development environment in Jupyter.\n",
    "\n",
    "Our next step is two fold. We want to be able to leave code in our notebooks that gets omitted when imported and we want to convert our notebooks into a package.\n",
    "\n",
    "## Converting notebooks to pure python packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "import nbformat\n",
    "from IPython import get_ipython"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_all_magics(string):\n",
    "    ret = ''\n",
    "    if string.startswith('%%'):\n",
    "        return ret\n",
    "    for line in string.splitlines():\n",
    "        if not line.lstrip().startswith('%'):\n",
    "            ret += line + '\\n'\n",
    "    return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_nb(filename):\n",
    "    return filename.endswith('.ipynb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_nb_source(filename, remove_magics=True):\n",
    "    assert is_nb(filename), 'File %s is not a Jupyter notebook' % filename\n",
    "    try:\n",
    "        nb = nbformat.read(filename, as_version=4)\n",
    "        code_cells = [c for c in nb['cells'] if c['cell_type'] == 'code']\n",
    "        nb_source = ''\n",
    "        for code_cell in code_cells:\n",
    "            if remove_magics:\n",
    "                nb_source += remove_all_magics(code_cell.source) + '\\n'\n",
    "            else: # turn magics into executable code\n",
    "                nb_source += shell.input_transformer_manager.transform_cell(code_cell.source)\n",
    "        return nb_source\n",
    "    except:\n",
    "        raise Exception('Notebook %s not readable' % filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nb_to_py(nbFile, pyPath=None, overwrite=False, remove_magics=True):\n",
    "    nbFile = os.path.expanduser(nbFile)\n",
    "    dirname, basename = os.path.split(nbFile)\n",
    "    if pyPath:\n",
    "        dirname = os.path.expanduser(pyPath)\n",
    "    basename = basename.split('.')[0] + '.py'\n",
    "    pyFile = os.path.join(dirname, basename)\n",
    "    if not overwrite:\n",
    "        assert not os.path.exists(pyFile), 'File at %s already exists. Use argument overwrite=True to overwrite' % pyFile\n",
    "    source = get_nb_source(nbFile, remove_magics)\n",
    "    with open(pyFile, \"w\") as pyFile:\n",
    "        pyFile.write(source)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nbs_to_pkg(nbDir, pyDir=None, remove_magics=True):\n",
    "    nbDir = os.path.expanduser(nbDir)\n",
    "    packageName = [dirname for dirname in nbDir.split('/') if dirname != ''][-1]\n",
    "    if pyDir:\n",
    "        pyDir = os.path.expanduser(pyDir)\n",
    "    else: \n",
    "        pyDir = nbDir\n",
    "    recurse(nbDir, pyDir, '', remove_magics)\n",
    "    \n",
    "def recurse(nbDir, pyDir, currentDir, remove_magics):\n",
    "    nbPath = os.path.join(nbDir, currentDir)\n",
    "    pyPath = os.path.join(pyDir, currentDir)\n",
    "    nbs = [child for child in os.listdir(nbPath) if not child.startswith('.') and child.endswith('.ipynb') and child != 'application.ipynb']\n",
    "    for nb in nbs:\n",
    "        nbFile = os.path.join(nbPath, nb)\n",
    "        os.makedirs(pyPath, exist_ok=True) \n",
    "        nb_to_py(nbFile, pyPath, overwrite=True)\n",
    "        \n",
    "    initFile = os.path.join(pyDir, '__init__.py')\n",
    "    with open(initFile, 'w') as f:\n",
    "        pass\n",
    "        \n",
    "    childDirs = [child for child in os.listdir(nbPath) if not child.startswith('.') and not child.startswith('__')]\n",
    "    childDirs = [child for child in childDirs if os.path.isdir(os.path.join(nbPath, child))]\n",
    "    if childDirs:\n",
    "        for childDir in childDirs:\n",
    "            newCurrent = os.path.join(currentDir, childDir)\n",
    "            recurse(nbDir, pyDir, newCurrent, remove_magics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "@magics_class\n",
    "class NotebookWriter(Magics):\n",
    "    \n",
    "    shell = get_ipython()\n",
    "    \n",
    "    @cell_magic\n",
    "    @magic_arguments.magic_arguments()\n",
    "    @magic_arguments.argument('--verbose', '-v',\n",
    "          help='Whether to omit cell from being written into python file.'\n",
    "    )\n",
    "    def omit(self, line='', cell=None):\n",
    "        args = magic_arguments.parse_argstring(self.omit, line)\n",
    "        self.shell.run_cell(cell)\n",
    "        \n",
    "    @line_magic\n",
    "    @magic_arguments.magic_arguments()\n",
    "    @magic_arguments.argument('--verbose', '-v',\n",
    "        help='Clear output and save the current notebook and write it to python.')\n",
    "    @magic_arguments.argument('notebook')\n",
    "    @magic_arguments.argument('--pyPath', '-d', default=os.getcwd())\n",
    "    @magic_arguments.argument('--nbPath', '-s', default=os.getcwd())\n",
    "    @magic_arguments.argument('--overwrite', '-f', default=False)\n",
    "    def nb_to_py(self, line='', cell=None):\n",
    "        args = magic_arguments.parse_argstring(self.nb_to_py, line)\n",
    "        nbFile = os.path.join(args.nbPath, args.notebook)\n",
    "        assert os.path.isfile(nbFile), 'Notebook not found at %s' % nbFile\n",
    "        assert os.path.isdir(args.pyPath), '%s is not an existing directory' % agrs.pyPath\n",
    "        nb_to_py(nbFile=nbFile, pyPath=args.pyPath, overwrite=args.overwrite)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "here = Path('.').resolve()\n",
    "nbs = here / 'nbs'\n",
    "pkg  = here / 'pkg'\n",
    "nbs_to_pkg(nbs/pkg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__init__.py  __pycache__/  another_module.py  child/  module.py\r\n"
     ]
    }
   ],
   "source": [
    "%ls pkg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Magic helpers\n",
    "\n",
    "The cell magic allows a cell to be executed without being included in the resulting python package. The line magic allows us to run nb_to_py on the current notebook, which could be used standalone, or within a package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core import magic_arguments\n",
    "from IPython.core.magic import line_magic, cell_magic, line_cell_magic, Magics, magics_class\n",
    "from IPython import get_ipython\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "@magics_class\n",
    "class NotebookWriter(Magics):\n",
    "    \n",
    "    shell = get_ipython()\n",
    "    \n",
    "    @cell_magic\n",
    "    @magic_arguments.magic_arguments()\n",
    "    @magic_arguments.argument('--verbose', '-v',\n",
    "          help='Whether to omit cell from being written into python file.'\n",
    "    )\n",
    "    def omit(self, line='', cell=None):\n",
    "        args = magic_arguments.parse_argstring(self.omit, line)\n",
    "        self.shell.run_cell(cell)\n",
    "        \n",
    "    @line_magic\n",
    "    @magic_arguments.magic_arguments()\n",
    "    @magic_arguments.argument('--verbose', '-v',\n",
    "        help='Clear output and save the current notebook and write it to python.')\n",
    "    @magic_arguments.argument('notebook')\n",
    "    @magic_arguments.argument('--pyPath', '-d', default=os.getcwd())\n",
    "    @magic_arguments.argument('--nbPath', '-s', default=os.getcwd())\n",
    "    @magic_arguments.argument('--overwrite', '-f', default=False)\n",
    "    def nb_to_py(self, line='', cell=None):\n",
    "        args = magic_arguments.parse_argstring(self.nb_to_py, line)\n",
    "        nbFile = os.path.join(args.nbPath, args.notebook)\n",
    "        assert os.path.isfile(nbFile), 'Notebook not found at %s' % nbFile\n",
    "        assert os.path.isdir(args.pyPath), '%s is not an existing directory' % agrs.pyPath\n",
    "        nb_to_py(nbFile=nbFile, pyPath=args.pyPath, overwrite=args.overwrite)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "shell = get_ipython()\n",
    "shell.register_magics(NotebookWriter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Code that is run but not included in the resulting py file\n"
     ]
    }
   ],
   "source": [
    "%%omit\n",
    "\n",
    "print('Code that is run but not included in the resulting py file')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "%nb_to_py  notebook_importer.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "import sys\r\n",
      "sys.path\r\n",
      "\r\n",
      "sys.meta_path\r\n",
      "\r\n",
      "import importlib\r\n",
      "from IPython import get_ipython\r\n",
      "import nbformat\r\n",
      "import os, sys\r\n",
      "from IPython.core.interactiveshell import InteractiveShell\r\n",
      "\r\n",
      "class NotebookLoader(importlib.abc.InspectLoader, importlib.abc.MetaPathFinder):\r\n",
      "    '''An abstract base class for a loader which implements the optional\r\n",
      "    PEP 302 protocol for loaders that inspect modules.\r\n",
      "    \r\n",
      "    Parameters:\r\n",
      "        path: path to the directory containing the module [string]\r\n",
      "        module_name: a name to use for the module. Path basename is default [string]\r\n",
      "        remove_magics: remove all cell and line magics from Notebook cell source code\r\n",
      "    '''\r\n",
      "    def __init__(self, module_path, remove_magics=True):\r\n",
      "        self.remove_magics = remove_magics\r\n",
      "        self.shell = InteractiveShell.instance()\r\n",
      "        self.module_path = os.path.expanduser(module_path)  \r\n",
      "        if not os.path.isdir(self.module_path):\r\n",
      "            raise FileNotFoundError('Module directory %s not found.' % self.module_path)\r",
      "\r\n",
      "        self.module_name = [dirname for dirname in self.module_path.split('/') if dirname != ''][-1]\r\n",
      "\r\n",
      "    def get_filename(self, fullname):\r\n",
      "        '''An abstract method that is to return the value of __file__ for the specified module.\r\n",
      "        If no path is available, ImportError is raised. If source code is available, then the \r\n",
      "        method should return the path to the source file\r\n",
      "        '''\r\n",
      "        parts = fullname.split('.')\r\n",
      "        if parts[0] != self.module_name:\r\n",
      "            raise ImportError('Module %s doesn\\'t exist' % parts[0])\r\n",
      "        if len(parts) == 1:\r\n",
      "            nb_path = 'this file does not exist'\r\n",
      "            init_path = '__init__.py'\r\n",
      "        elif len(parts) == 2:\r\n",
      "            nb_path = parts[1] + '.ipynb'\r\n",
      "            init_path = os.path.join(parts[1], '__init__.py')\r\n",
      "        else:\r\n",
      "            parts[-1] = parts[-1] + '.ipynb'\r\n",
      "            nb_path = os.path.join(parts[1], *parts[2:])\r\n",
      "            init_path = os.path.join(parts[1], *parts[2:], '__init__.py')\r\n",
      "            \r\n",
      "        for path in (nb_path, init_path):\r\n",
      "            fullpath = os.path.join(self.module_path, path)\r\n",
      "            if os.path.exists(fullpath):\r\n",
      "                return fullpath\r\n",
      "        raise ImportError('Module %s does not exist at %s' % (fullname, fullpath))\r\n",
      "            \r\n",
      "        \r\n",
      "    def get_code(self, fullname):\r\n",
      "        '''Return the code object for a module, or None if the module does not have a code\r\n",
      "        object (as would be the case, for example, for a built-in module). Raise an \r\n",
      "        ImportError if loader cannot find the requested module.\r\n",
      "        '''\r\n",
      "        filename = self.get_filename(fullname)\r\n",
      "        return self.get_code_from_file(filename)\r\n",
      "\r\n",
      "    def get_source(self, fullname):\r\n",
      "        '''An abstract method to return the source of a module. It is returned as a\r\n",
      "        text string using universal newlines, translating all recognized line separators\r\n",
      "        into '\\n' characters. Returns None if no source is available (e.g. a built-in module).\r\n",
      "        Raises ImportError if the loader cannot find the module specified.\r\n",
      "        '''\r\n",
      "        return self.get_code(fullname)\r\n",
      "       \r\n",
      "    def get_code_from_file(self, filename):\r\n",
      "        if self.is_nb(filename):\r\n",
      "            return self.get_nb_source(filename)\r\n",
      "        else:\r\n",
      "            return self.get_init_source(filename)\r\n",
      "        \r\n",
      "            \r\n",
      "    def get_init_source(self, filename):\r\n",
      "        with open(filename, 'r') as f:\r\n",
      "            return f.read()\r\n",
      "        raise ModuleNotFoundError('Filename %s not readable' % filename)\r\n",
      "        \r\n",
      "    def get_nb_source(self, filename):\r\n",
      "        try:\r\n",
      "            nb = nbformat.read(filename, as_version=4)\r\n",
      "            code_cells = [c for c in nb['cells'] if c['cell_type'] == 'code']\r\n",
      "            nb_source = ''\r\n",
      "            for code_cell in code_cells:\r\n",
      "                if self.remove_magics:\r\n",
      "                    nb_source += self.remove_all_magics(code_cell.source) + '\\n'\r\n",
      "                else: # turn magics into executable code\r\n",
      "                    nb_source += self.shell.input_transformer_manager.transform_cell(code_cell.source)\r\n",
      "            return nb_source\r\n",
      "        except:\r\n",
      "            raise ModuleNotFoundExpection('Notebook %s not readable' % filename)\r\n",
      "\r\n",
      "    def remove_all_magics(self, string):\r\n",
      "        ret = ''\r\n",
      "        if string.startswith('%%'):\r\n",
      "            return ret\r\n",
      "        for line in string.splitlines():\r\n",
      "            if not line.lstrip().startswith('%'):\r\n",
      "                ret += line + '\\n'\r\n",
      "        return ret\r\n",
      "    \r\n",
      "    def is_nb(self, filename):\r\n",
      "        return filename.endswith('.ipynb')\r\n",
      "\r\n",
      "    def is_package(self, fullname):\r\n",
      "        '''An abstract method to return a true value if the module is a package, a false\r\n",
      "        value otherwise. ImportError is raised if the loader cannot find the module.\r\n",
      "        '''\r\n",
      "        return fullname == self.module_name\r\n",
      "    \r\n",
      "    def exec_module(self, module):\r\n",
      "        '''An abstract method that executes the module in its own namespace when a module \r\n",
      "        is imported or reloaded. The module should already be initialized when exec_module() \r\n",
      "        is called. When this method exists, create_module() must be defined.\r\n",
      "        '''\r\n",
      "        try:\r\n",
      "            exec(self.get_code_from_file(module.__file__), module.__dict__)\r\n",
      "        except:\r\n",
      "            print('Error executing %s' % module.__file__)\r\n",
      "            print(self.get_code_from_file(module.__file__))\r\n",
      "            raise\r\n",
      "\r\n",
      "class NotebookFinder(importlib.abc.MetaPathFinder):\r\n",
      "    \r\n",
      "    def __init__(self, module_path):\r\n",
      "        self.module_path = module_path\r\n",
      "\r\n",
      "    def find_spec(self, fullname, path, target=None):\r\n",
      "        '''An abstract method for finding a spec for the specified module. If this\r\n",
      "        is a top-level import, path will be None. Otherwise, this is a search for\r\n",
      "        a subpackage or module and path will be the value of __path__ from the parent\r\n",
      "        package. If a spec cannot be found, None is returned. When passed in, target\r\n",
      "        is a module object that the finder may use to make a more educated guess about\r\n",
      "        what spec to return. importlib.util.spec_from_loader() may be useful for \r\n",
      "        implementing concrete MetaPathFinders.\r\n",
      "        '''\r\n",
      "        loader = NotebookLoader(self.module_path) # __loader__: The Loader that should be used when loading the module. Finders should always set this.\r\n",
      "        name = fullname # __name__: A string for the fully-qualified name of the module.\r\n",
      "        # __cached__: String for where the compiled module should be stored (or None).\r\n",
      "        # __package__: (Read-only) The fully-qualified name of the package under which the module should be loaded as a submodule (or the empty string for top-level modules). For packages, it is the same as __name__.\r\n",
      "        location = loader.get_filename(fullname) # __file__: The path to where the module data is stored (not set for built-in modules).\r\n",
      "        self.submodule_search_locations = None # __path__: List of strings for where to find submodules, if a package (None otherwise).\r\n",
      "        if not location.endswith('.ipynb'):\r\n",
      "            (spec_dir, _) = os.path.split(location)\r\n",
      "            self.submodule_search_locations = []\r\n",
      "            self.submodule_search(spec_dir)\r\n",
      "        spec = importlib.util.spec_from_file_location(name, location, loader=loader, submodule_search_locations=self.submodule_search_locations)\r\n",
      "        return spec\r\n",
      "    \r\n",
      "    def submodule_search(self, location):\r\n",
      "        locations = [os.path.join(location, f) for f in os.listdir(location) if not f.startswith('.')]\r\n",
      "        locations = [d for d in locations if os.path.isdir(d) and os.path.isfile(os.path.join(d, '__init__.py'))]\r\n",
      "        self.submodule_search_locations = locations\r\n",
      "\r\n",
      "sys.meta_path.append(NotebookFinder('nbs'))\r\n",
      "\r\n",
      "print(sys.meta_path)\r\n",
      "\r\n",
      "\r\n",
      "from nbs.module import hello_world\r\n",
      "\r\n",
      "hello_world()\r\n",
      "\r\n",
      "from nbs.another_module import conversation\r\n",
      "\r\n",
      "conversation()\r\n",
      "\r\n",
      "\r\n",
      "from nbs.child.baby_module import babble\r\n",
      "\r\n",
      "babble()\r\n",
      "\r\n",
      "import os, sys\r\n",
      "import nbformat\r\n",
      "from IPython import get_ipython\r\n",
      "\r\n",
      "def remove_all_magics(string):\r\n",
      "    ret = ''\r\n",
      "    if string.startswith('%%'):\r\n",
      "        return ret\r\n",
      "    for line in string.splitlines():\r\n",
      "        if not line.lstrip().startswith('%'):\r\n",
      "            ret += line + '\\n'\r\n",
      "    return ret\r\n",
      "\r\n",
      "def is_nb(filename):\r\n",
      "    return filename.endswith('.ipynb')\r\n",
      "\r\n",
      "def get_nb_source(filename, remove_magics=True):\r\n",
      "    assert is_nb(filename), 'File %s is not a Jupyter notebook' % filename\r\n",
      "    try:\r\n",
      "        nb = nbformat.read(filename, as_version=4)\r\n",
      "        code_cells = [c for c in nb['cells'] if c['cell_type'] == 'code']\r\n",
      "        nb_source = ''\r\n",
      "        for code_cell in code_cells:\r\n",
      "            if remove_magics:\r\n",
      "                nb_source += remove_all_magics(code_cell.source) + '\\n'\r\n",
      "            else: # turn magics into executable code\r\n",
      "                nb_source += shell.input_transformer_manager.transform_cell(code_cell.source)\r\n",
      "        return nb_source\r\n",
      "    except:\r\n",
      "        raise Exception('Notebook %s not readable' % filename)\r\n",
      "\r\n",
      "def nb_to_py(nbFile, pyPath=None, overwrite=False, remove_magics=True):\r\n",
      "    nbFile = os.path.expanduser(nbFile)\r\n",
      "    dirname, basename = os.path.split(nbFile)\r\n",
      "    if pyPath:\r\n",
      "        dirname = os.path.expanduser(pyPath)\r\n",
      "    basename = basename.split('.')[0] + '.py'\r\n",
      "    pyFile = os.path.join(dirname, basename)\r\n",
      "    if not overwrite:\r\n",
      "        assert not os.path.exists(pyFile), 'File at %s already exists. Use argument overwrite=True to overwrite' % pyFile\r\n",
      "    source = get_nb_source(nbFile, remove_magics)\r\n",
      "    with open(pyFile, \"w\") as pyFile:\r\n",
      "        pyFile.write(source)\r\n",
      "\r\n",
      "def nbs_to_pkg(nbDir, pyDir=None, remove_magics=True):\r\n",
      "    nbDir = os.path.expanduser(nbDir)\r\n",
      "    packageName = [dirname for dirname in nbDir.split('/') if dirname != ''][-1]\r\n",
      "    if pyDir:\r\n",
      "        pyDir = os.path.expanduser(pyDir)\r\n",
      "    else: \r\n",
      "        pyDir = nbDir\r\n",
      "    recurse(nbDir, pyDir, '', remove_magics)\r\n",
      "    \r\n",
      "def recurse(nbDir, pyDir, currentDir, remove_magics):\r\n",
      "    nbPath = os.path.join(nbDir, currentDir)\r\n",
      "    pyPath = os.path.join(pyDir, currentDir)\r\n",
      "    nbs = [child for child in os.listdir(nbPath) if not child.startswith('.') and child.endswith('.ipynb') and child != 'application.ipynb']\r\n",
      "    for nb in nbs:\r\n",
      "        nbFile = os.path.join(nbPath, nb)\r\n",
      "        os.makedirs(pyPath, exist_ok=True) \r\n",
      "        nb_to_py(nbFile, pyPath, overwrite=True)\r\n",
      "        \r\n",
      "    initFile = os.path.join(pyDir, '__init__.py')\r\n",
      "    with open(initFile, 'w') as f:\r\n",
      "        pass\r\n",
      "        \r\n",
      "    childDirs = [child for child in os.listdir(nbPath) if not child.startswith('.') and not child.startswith('__')]\r\n",
      "    childDirs = [child for child in childDirs if os.path.isdir(os.path.join(nbPath, child))]\r\n",
      "    if childDirs:\r\n",
      "        for childDir in childDirs:\r\n",
      "            newCurrent = os.path.join(currentDir, childDir)\r\n",
      "            recurse(nbDir, pyDir, newCurrent, remove_magics)\r\n",
      "\r\n",
      "@magics_class\r\n",
      "class NotebookWriter(Magics):\r\n",
      "    \r\n",
      "    shell = get_ipython()\r\n",
      "    \r\n",
      "    @cell_magic\r\n",
      "    @magic_arguments.magic_arguments()\r\n",
      "    @magic_arguments.argument('--verbose', '-v',\r\n",
      "          help='Whether to omit cell from being written into python file.'\r\n",
      "    )\r\n",
      "    def omit(self, line='', cell=None):\r\n",
      "        args = magic_arguments.parse_argstring(self.omit, line)\r\n",
      "        self.shell.run_cell(cell)\r\n",
      "        \r\n",
      "    @line_magic\r\n",
      "    @magic_arguments.magic_arguments()\r\n",
      "    @magic_arguments.argument('--verbose', '-v',\r\n",
      "        help='Clear output and save the current notebook and write it to python.')\r\n",
      "    @magic_arguments.argument('notebook')\r\n",
      "    @magic_arguments.argument('--pyPath', '-d', default=os.getcwd())\r\n",
      "    @magic_arguments.argument('--nbPath', '-s', default=os.getcwd())\r\n",
      "    @magic_arguments.argument('--overwrite', '-f', default=False)\r\n",
      "    def nb_to_py(self, line='', cell=None):\r\n",
      "        args = magic_arguments.parse_argstring(self.nb_to_py, line)\r\n",
      "        nbFile = os.path.join(args.nbPath, args.notebook)\r\n",
      "        assert os.path.isfile(nbFile), 'Notebook not found at %s' % nbFile\r\n",
      "        assert os.path.isdir(args.pyPath), '%s is not an existing directory' % agrs.pyPath\r\n",
      "        nb_to_py(nbFile=nbFile, pyPath=args.pyPath, overwrite=args.overwrite)\r\n",
      "\r\n",
      "here = Path('.').resolve()\r\n",
      "nbs = here / 'nbs'\r\n",
      "pkg  = here / 'pkg'\r\n",
      "nbs_to_pkg(nbs/pkg)\r\n",
      "\r\n",
      "\r\n",
      "from IPython.core import magic_arguments\r\n",
      "from IPython.core.magic import line_magic, cell_magic, line_cell_magic, Magics, magics_class\r\n",
      "from IPython import get_ipython\r\n",
      "import os\r\n",
      "\r\n",
      "@magics_class\r\n",
      "class NotebookWriter(Magics):\r\n",
      "    \r\n",
      "    shell = get_ipython()\r\n",
      "    \r\n",
      "    @cell_magic\r\n",
      "    @magic_arguments.magic_arguments()\r\n",
      "    @magic_arguments.argument('--verbose', '-v',\r\n",
      "          help='Whether to omit cell from being written into python file.'\r\n",
      "    )\r\n",
      "    def omit(self, line='', cell=None):\r\n",
      "        args = magic_arguments.parse_argstring(self.omit, line)\r\n",
      "        self.shell.run_cell(cell)\r\n",
      "        \r\n",
      "    @line_magic\r\n",
      "    @magic_arguments.magic_arguments()\r\n",
      "    @magic_arguments.argument('--verbose', '-v',\r\n",
      "        help='Clear output and save the current notebook and write it to python.')\r\n",
      "    @magic_arguments.argument('notebook')\r\n",
      "    @magic_arguments.argument('--pyPath', '-d', default=os.getcwd())\r\n",
      "    @magic_arguments.argument('--nbPath', '-s', default=os.getcwd())\r\n",
      "    @magic_arguments.argument('--overwrite', '-f', default=False)\r\n",
      "    def nb_to_py(self, line='', cell=None):\r\n",
      "        args = magic_arguments.parse_argstring(self.nb_to_py, line)\r\n",
      "        nbFile = os.path.join(args.nbPath, args.notebook)\r\n",
      "        assert os.path.isfile(nbFile), 'Notebook not found at %s' % nbFile\r\n",
      "        assert os.path.isdir(args.pyPath), '%s is not an existing directory' % agrs.pyPath\r\n",
      "        nb_to_py(nbFile=nbFile, pyPath=args.pyPath, overwrite=args.overwrite)\r\n",
      "\r\n",
      "shell = get_ipython()\r\n",
      "shell.register_magics(NotebookWriter)\r\n",
      "\r\n",
      "\r\n",
      "\r\n"
     ]
    }
   ],
   "source": [
    "%cat notebook_importer.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}