Merge pull request #252 from mmerickel/feature.pshell-setup

Feature.pshell setup

3 files changed, 150 insertions, 12 deletions

diff --git a/docs/narr/commandline.rst b/docs/narr/commandline.rst
index 1c9d0b15c..bb004c446 100644
--- a/docs/narr/commandline.rst
+++ b/docs/narr/commandline.rst
@@ -156,6 +156,7 @@ name ``MyProject`` as a section name:
       request      Active request object.
       root         Root of the default resource tree.
       root_factory Default root factory used to create `root`.
+
     >>> root
     <myproject.resources.MyResource object at 0x445270>
     >>> registry
@@ -191,12 +192,18 @@ Press ``Ctrl-D`` to exit the interactive shell (or ``Ctrl-Z`` on Windows).
 Extending the Shell
 ~~~~~~~~~~~~~~~~~~~
 
-It is sometimes convenient when using the interactive shell often to have
-some variables significant to your application already loaded as globals when
+It is convenient when using the interactive shell often to have some
+variables significant to your application already loaded as globals when
 you start the ``pshell``. To facilitate this, ``pshell`` will look for a
 special ``[pshell]`` section in your INI file and expose the subsequent
 key/value pairs to the shell.  Each key is a variable name that will be
 global within the pshell session; each value is a :term:`dotted Python name`.
+If specified, the special key ``setup`` should be a :term:`dotted Python name`
+pointing to a callable that accepts the dictionary of globals that will
+be loaded into the shell. This allows for some custom initializing code
+to be executed each time the ``pshell`` is run. The ``setup`` callable
+can also be specified from the commandline using the ``--setup`` option
+which will override the key in the INI file.
 
 For example, you want to expose your model to the shell, along with the
 database session so that you can mutate the model on an actual database.
@@ -206,12 +213,33 @@ Here, we'll assume your model is stored in the ``myapp.models`` package.
    :linenos:
 
    [pshell]
+   setup = myapp.lib.pshell.setup
    m = myapp.models
    session = myapp.models.DBSession
    t = transaction
 
+By defining the ``setup`` callable, we will create the module
+``myapp.lib.pshell`` containing a callable named ``setup`` that will receive
+the global environment before it is exposed to the shell. Here we mutate the
+environment's request as well as add a new value containing a WebTest version
+of the application to which we can easily submit requests.
+
+.. code-block:: python
+    :linenos:
+
+    # myapp/lib/pshell.py
+    from webtest import TestApp
+
+    def setup(env):
+        env['request'].host = 'www.example.com'
+        env['request'].scheme = 'https'
+        env['testapp'] = TestApp(env['app'])
+
 When this INI file is loaded, the extra variables ``m``, ``session`` and
-``t`` will be available for use immediately. For example:
+``t`` will be available for use immediately. Since a ``setup`` callable
+was also specified, it is executed and a new variable ``testapp`` is
+exposed, and the request is configured to generate urls from the host
+``http://www.example.com``. For example:
 
 .. code-block:: text
 
@@ -226,12 +254,17 @@ When this INI file is loaded, the extra variables ``m``, ``session`` and
       request      Active request object.
       root         Root of the default resource tree.
       root_factory Default root factory used to create `root`.
+      testapp      <webtest.TestApp object at ...>
 
     Custom Variables:
       m            myapp.models
       session      myapp.models.DBSession
       t            transaction
-    >>>
+
+    >>> testapp.get('/')
+    <200 OK text/html body='<!DOCTYPE...l>\n'/3337>
+    >>> request.route_url('home')
+    'https://www.example.com/'
 
 .. index::
    single: IPython
diff --git a/pyramid/paster.py b/pyramid/paster.py
index f9765a07c..ec55dc126 100644
--- a/pyramid/paster.py
+++ b/pyramid/paster.py
@@ -129,22 +129,37 @@ class PShellCommand(PCommand):
                       action='store_true',
                       dest='disable_ipython',
                       help="Don't use IPython even if it is available")
+    parser.add_option('--setup',
+                      dest='setup',
+                      help=("A callable that will be passed the environment "
+                            "before it is made available to the shell. This "
+                            "option will override the 'setup' key in the "
+                            "[pshell] ini section."))
 
     ConfigParser = ConfigParser.ConfigParser # testing
 
+    loaded_objects = {}
+    object_help = {}
+    setup = None
+
     def pshell_file_config(self, filename):
-        resolver = DottedNameResolver(None)
-        self.loaded_objects = {}
-        self.object_help = {}
         config = self.ConfigParser()
         config.read(filename)
         try:
             items = config.items('pshell')
         except ConfigParser.NoSectionError:
             return
+
+        resolver = DottedNameResolver(None)
+        self.loaded_objects = {}
+        self.object_help = {}
+        self.setup = None
         for k, v in items:
-            self.loaded_objects[k] = resolver.maybe_resolve(v)
-            self.object_help[k] = v
+            if k == 'setup':
+                self.setup = v
+            else:
+                self.loaded_objects[k] = resolver.maybe_resolve(v)
+                self.object_help[k] = v
 
     def command(self, shell=None):
         config_uri = self.args[0]
@@ -167,6 +182,24 @@ class PShellCommand(PCommand):
         env_help['root_factory'] = (
             'Default root factory used to create `root`.')
 
+        # override use_script with command-line options
+        if self.options.setup:
+            self.setup = self.options.setup
+
+        if self.setup:
+            # store the env before muddling it with the script
+            orig_env = env.copy()
+
+            # call the setup callable
+            resolver = DottedNameResolver(None)
+            setup = resolver.maybe_resolve(self.setup)
+            setup(env)
+
+            # remove any objects from default help that were overidden
+            for k, v in env.iteritems():
+                if k not in orig_env or env[k] != orig_env[k]:
+                    env_help[k] = v
+
         # load the pshell section of the ini file
         env.update(self.loaded_objects)
 
@@ -176,7 +209,7 @@ class PShellCommand(PCommand):
                 del env_help[k]
 
         # generate help text
-        help = '\n'
+        help = ''
         if env_help:
             help += 'Environment:'
             for var in sorted(env_help.keys()):
@@ -204,7 +237,7 @@ class PShellCommand(PCommand):
         def shell(env, help):
             cprt = 'Type "help" for more information.'
             banner = "Python %s on %s\n%s" % (sys.version, sys.platform, cprt)
-            banner += '\n' + help + '\n'
+            banner += '\n\n' + help + '\n'
             interact(banner, local=env)
         return shell
 
@@ -217,7 +250,7 @@ class PShellCommand(PCommand):
             except ImportError:
                 return None
         def shell(env, help):
-            IPShell = IPShellFactory(banner2=help, user_ns=env)
+            IPShell = IPShellFactory(banner2=help + '\n', user_ns=env)
             IPShell()
         return shell
 
diff --git a/pyramid/tests/test_paster.py b/pyramid/tests/test_paster.py
index 36c3a51be..995e05d46 100644
--- a/pyramid/tests/test_paster.py
+++ b/pyramid/tests/test_paster.py
@@ -21,6 +21,7 @@ class TestPShellCommand(unittest.TestCase):
             class Options(object): pass
             self.options = Options()
             self.options.disable_ipython = True
+            self.options.setup = None
             cmd.options = self.options
         return cmd
 
@@ -156,6 +157,77 @@ class TestPShellCommand(unittest.TestCase):
         self.assertTrue(self.bootstrap.closer.called)
         self.assertTrue(shell.help)
 
+    def test_command_setup(self):
+        command = self._makeOne()
+        def setup(env):
+            env['a'] = 1
+            env['root'] = 'root override'
+        self.config_factory.items = [('setup', setup)]
+        shell = DummyShell()
+        command.command(shell)
+        self.assertTrue(self.config_factory.parser)
+        self.assertEqual(self.config_factory.parser.filename,
+                         '/foo/bar/myapp.ini')
+        self.assertEqual(self.bootstrap.a[0], '/foo/bar/myapp.ini#myapp')
+        self.assertEqual(shell.env, {
+            'app':self.bootstrap.app, 'root':'root override',
+            'registry':self.bootstrap.registry,
+            'request':self.bootstrap.request,
+            'root_factory':self.bootstrap.root_factory,
+            'a':1,
+        })
+        self.assertTrue(self.bootstrap.closer.called)
+        self.assertTrue(shell.help)
+
+    def test_command_loads_check_variable_override_order(self):
+        command = self._makeOne()
+        model = Dummy()
+        def setup(env):
+            env['a'] = 1
+            env['m'] = 'model override'
+            env['root'] = 'root override'
+        self.config_factory.items = [('setup', setup), ('m', model)]
+        shell = DummyShell()
+        command.command(shell)
+        self.assertTrue(self.config_factory.parser)
+        self.assertEqual(self.config_factory.parser.filename,
+                         '/foo/bar/myapp.ini')
+        self.assertEqual(self.bootstrap.a[0], '/foo/bar/myapp.ini#myapp')
+        self.assertEqual(shell.env, {
+            'app':self.bootstrap.app, 'root':'root override',
+            'registry':self.bootstrap.registry,
+            'request':self.bootstrap.request,
+            'root_factory':self.bootstrap.root_factory,
+            'a':1, 'm':model,
+        })
+        self.assertTrue(self.bootstrap.closer.called)
+        self.assertTrue(shell.help)
+
+    def test_command_loads_setup_from_options(self):
+        command = self._makeOne()
+        def setup(env):
+            env['a'] = 1
+            env['root'] = 'root override'
+        model = Dummy()
+        self.config_factory.items = [('setup', 'abc'),
+                                     ('m', model)]
+        command.options.setup = setup
+        shell = DummyShell()
+        command.command(shell)
+        self.assertTrue(self.config_factory.parser)
+        self.assertEqual(self.config_factory.parser.filename,
+                         '/foo/bar/myapp.ini')
+        self.assertEqual(self.bootstrap.a[0], '/foo/bar/myapp.ini#myapp')
+        self.assertEqual(shell.env, {
+            'app':self.bootstrap.app, 'root':'root override',
+            'registry':self.bootstrap.registry,
+            'request':self.bootstrap.request,
+            'root_factory':self.bootstrap.root_factory,
+            'a':1, 'm':model,
+        })
+        self.assertTrue(self.bootstrap.closer.called)
+        self.assertTrue(shell.help)
+
     def test_command_custom_section_override(self):
         command = self._makeOne()
         dummy = Dummy()