Source

pythonhk / src / pythonhk / console.py

Full commit
import os
import sys
from argparse import ArgumentParser
from code import InteractiveConsole

import cherrypy
from sqlalchemy.engine import engine_from_config

class Console(InteractiveConsole):

    def __init__(self, confs, prompt):

        if confs:
            if isinstance(confs, (tuple, list)):
                confs = list(confs)
            else:
                confs = [confs]

            config = {}
            for conf in confs:
                cherrypy._cpconfig.merge(config, conf)
            self.config = config
        else:
            self.config = None

        try:
            import readline
        except ImportError, e:
            print e
        else:
            import rlcompleter

        startupfile = os.environ.get("PYTHONSTARTUP")
        if startupfile: execfile(startupfile, {}, {})

        sys.ps1 = "[%s]>>> " % prompt
        sys.ps2 = "[%s]... " % prompt

        self.prompt = prompt

        InteractiveConsole.__init__(self, locals=self.get_locals())

    def get_locals(self):
        import sys, os, os.path, time, datetime, pprint, inspect
        import sqlalchemy
        import sqlalchemy.orm
        import cherrypy
        from pythonhk import models

        lcls = dict(locals())

        if getattr(models, '__all__'):
            for name in models.__all__:
                lcls[name] = getattr(models, name)
        else:
            for name, obj in vars(models).iteritems():
                if not name.startswith("_"):
                    lcls[name] = obj

        if self.config:
            section = self.config.get("sqlalchemy_engine")
            if section:
                engine = engine_from_config(section)
                Session = sqlalchemy.orm.sessionmaker()
                Session.configure(bind=engine)
                metadata = models.metadata
                metadata.bind = engine
                lcls['create_all'] = metadata.create_all
                lcls['drop_all'] = metadata.drop_all
                lcls['session'] = session = Session()
                import atexit
                atexit.register(session.close)

        return lcls

    def raw_input(self, *args, **kw):
        try:
            r = InteractiveConsole.raw_input(self, *args, **kw)
            for encoding in (getattr(sys.stdin, 'encoding', None),
                             sys.getdefaultencoding(), 'utf-8', 'latin-1'):
                if encoding:
                    try:
                        return r.decode(encoding)
                    except UnicodeError:
                        pass
                    return r
        except EOFError:
            self.write(os.linesep)
            session = self.locals.get("session")
            if session is not None and \
                getattr(session, "new", None) or \
                getattr(session, "dirty", None) or \
                getattr(session, "deleted", None):

                r = raw_input("Do you wish to commit your "
                              "database changes? [Y/n]")
                if not r.lower().startswith("n"):
                    self.push("session.flush()")
            raise

def main():

    argparser = ArgumentParser(description="Playground for prototyping SQLAlchemy queries against the models and experimenting with CherryPy's tools.")
    argparser.add_argument("-p", "--prompt", default="pythonhk")
    argparser.add_argument("-f", "--config-file", nargs=1, help="The configuration file to set up the SQLAlchemy session. If any.")

    args = argparser.parse_args()

    banner = """
*****************************************************************************
* If the configuration file you specified contains a [sqlalchemy_engine]    *
* section, a default SQLAlchemy engine and session should have been created *
* for you automatically already.                                            *
*****************************************************************************
"""

    Console(args.config_file, args.prompt).interact(banner)