Saturday morning hack: Automatically import your app's models when starting IPython

May 18, 2014 14:45 / python / 0 comments

I use Flask and peewee for all my personal projects, and wanted an easy way to automatically open an IPython shell with all my project's models in the namespace. If you've used the excellent django-extensions project, you may be familiar with the shell_plus command, which does the same thing. If your project has multiple models spread across several modules, this kind of hack can save you a lot of keystrokes.

Normally my workflow would look something like this:

$ ipython

In [1]: from blog.entries.blueprint import Entry, Comment

In [2]: from blog.models import Page

In [3]: # do something with these models.

Typing these import statements over and over gets pretty tedious. So here is the script I wrote that does it for me. It walks the directory finding non-executable Python files and automatically imports them, loading any Model subclasses into the namespace:

#!/usr/bin/env python

from importlib import import_module
import inspect
import os
from stat import ST_MODE
import sys

from IPython import embed
from peewee import Model


def main(site):
    namespace = {}

    for dirname, subdirs, filenames in os.walk(site or '.'):
        # Not a package, so skip this directory.
        if '__init__.py' not in filenames:
            continue

        # Create a dotted path representation of this directory.
        module_base = dirname.strip('.').strip('/').replace('/', '.')

        # Iterate over all the python files in this directory searching for
        # model subclasses.
        for filename in filter(lambda s: s.endswith('.py'), filenames):
            perms = os.stat(os.path.join(dirname, filename))[ST_MODE]

            # Skip over any executable Python scripts.
            if perms & 0100:
                continue

            # If this is the `__init__` in the root directory, skip over.
            if not module_base and filename == '__init__.py':
                continue

            if filename == '__init__.py':
                module_name = module_base
            elif module_base:
                module_name = '.'.join((module_base, filename[:-3]))
            else:
                module_name = filename[:-3]

            module = import_module(module_name)

            for name in dir(module):
                obj = getattr(module, name)
                if inspect.isclass(obj) and issubclass(obj, Model):
                    namespace[obj.__name__] = obj

    # Start up IPython with all the models in the namespace.
    embed(user_ns=namespace)

if __name__ == '__main__':
    if len(sys.argv) == 2:
        site = sys.argv[1]
    else:
        site = ''
    main(site)

I've put this script on my PATH so whenever I'm working on a Flask project, I can just run pyshell and everything just works!

Thanks for reading

Thanks for taking the time to read this post, I hope you found it helpful. As always I would appreciate any suggestions on how to improve this script -- feel free to leave a comment below.

If you're interested in more projects like this, check out the saturday-morning hack posts.

Comments (0)

Commenting has been closed.