#!/usr/bin/env python
# coding: utf-8

# # The ModuleMagic cell magic
# Gordon Bean, April 2015

# ## Motivation
# - %%file magic is nice for creating files - it is nice to see the contents of the file in the notebook.
# - %%cython allows you to create cython modules directly from notebook cells
# 
# I want to create python modules directly in the notebook. Like %%file, I want to be able to see the contents of the module, rather than having an additional editor for maintaining the file. 
# 
# However, rather than just using %%file, I want the code to be automatically imported into my namespace as a module, much like %%cython. New versions of the cell should override existing modules. Such functionality could be obtained using %%file, import, and importlib.reload, but would require two input cells. 
# 
# So, here I introduce the %%module cell magic. It saves the cell contents in a temporary file (unique to the underlying python kernel) and imports that file to the `__main__` namespace as a module.

# ## Loading the magic

# In[1]:


from beans.magics import ModuleMagics
ip = get_ipython()
ip.register_magics(ModuleMagics)


# ## Using the magic

# In[2]:


get_ipython().run_cell_magic('module', 'foobar', 'a = 1\nb = 2\n\ndef baz(c, d):\n    return c + d\n')


# In[3]:


foobar.a, foobar.b


# In[4]:


foobar.baz(4,5)


# ## Under the hood
# The ultimate look under the hood is to examine the source code.
# 
# Here is the "cliff-notes" version:
# 
# - Cell contents are saved in a kernel-specific directory in `/tmp/.tmp-modules/`. 
# - This module directory is added to `sys.path`, so they can be imported again by the same kernel (including in other modules created by %%module). 
# - The first word after %%module is taken as the name of the file, and therefore becomes the name of the module. 
#  - It should conform to python syntax requirements, but this is not enforced. As the module is added to sys.modules and the global namespace via `__dict__` interfaces, any name will probably work, but can't be accessed using normal syntax - I advise avoiding pathology.
# - The module is reloaded everytime the cell is run. Existing entries in sys.modules are deleted before the module is loaded again.
# - The module is loaded into the global namespace (i.e. `sys.modules['__main__'].__dict__`). The ModuleMagics class can be instantiated with the keyword `namespace` and a dictionary representing the namespace %%module modules will be loaded into. 

# ## Future directions
# - Add second argument to specify the package. 
# - Use ArgParse to handle magic command parsing.
# - Perhaps there is a better way to get a kernel-specific ID than parsing the connection_file path.

# ## The current source code
# This is likely to change, but this is the current version:

# In[ ]:


## beans.magics
# Gordon Bean, April 2015

# To use these magics, you must register them with IPython
# e.g:
# from beans.magics import ModuleMagics
# ip = get_ipython()
# ip.register_magics(ModuleMagics)

from IPython.core.magic import Magics, magics_class, cell_magic
import os, sys, importlib

@magics_class
class ModuleMagics(Magics):
    '''Magics for creating modules in IPython.'''
    
    def __init__(self, shell=None, namespace=None):
        if shell is None:
            shell = get_ipython()
            
        super(ModuleMagics, self).__init__(shell)
        
        self.namespace = namespace
        
        # Get the kernel id
        self.kernelID = os.path.basename(shell.kernel.config['IPKernelApp']['connection_file'])[:-5]
        
        # Create kernel-specific tmp-module directory
        self.module_dir = os.path.join('/tmp/.tmp-modules', self.kernelID)
        os.makedirs(self.module_dir, exist_ok=True)
    
    def __del__(self):
        # Remove module_dir from file system and sys.path
        # I'm not sure this works - evidence so far says no...
        tmpfiles = os.listdir(self.module_dir)
        for file in tmpfiles:
            os.remove(os.path.join(self.module_dir, file))
            
        os.rmdir(self.module_dir)
        
        sys.path.remove(self.module_dir)
    
    @cell_magic
    def module(self, line, cell):
        '''Import the cell as a module.'''
        # Parse module name
        tokens = line.split()
        name = tokens[0]

        # Save to file
        filename = os.path.join(self.module_dir, name + '.py')
        with open(filename, 'w') as f:
            f.write(cell)

        # Import module
        if self.module_dir not in sys.path:
            sys.path.insert(0, self.module_dir)

        namespace = self.namespace if self.namespace else sys.modules['__main__'].__dict__
        
        if name in namespace:
            # Always reload
            del namespace[name]

        module = importlib.import_module(name)
        namespace[name] = module

        return namespace[name]