Added support for 'import' key in [pshell] ini section.

3 files changed, 77 insertions, 5 deletions

diff --git a/pyramid/paster.py b/pyramid/paster.py
index f9765a07c..87abdddd4 100644
--- a/pyramid/paster.py
+++ b/pyramid/paster.py
@@ -132,19 +132,28 @@ class PShellCommand(PCommand):
 
     ConfigParser = ConfigParser.ConfigParser # testing
 
+    loaded_objects = {}
+    object_help = {}
+    use_script = 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.use_script = None
         for k, v in items:
-            self.loaded_objects[k] = resolver.maybe_resolve(v)
-            self.object_help[k] = v
+            if k == 'import':
+                self.use_script = 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 +176,19 @@ class PShellCommand(PCommand):
         env_help['root_factory'] = (
             'Default root factory used to create `root`.')
 
+        if self.use_script:
+            # store the env before muddling it with the script
+            orig_env = env.copy()
+
+            # do this instead of an eval() to respect __all__
+            exec 'from %s import *' % self.use_script in env
+            env.pop('__builtins__', None)
+
+            # remove any objects from default help that were overidden
+            for k, v in orig_env.iteritems():
+                if env[k] != orig_env[k]:
+                    del env_help[k]
+
         # load the pshell section of the ini file
         env.update(self.loaded_objects)
 
@@ -187,6 +209,9 @@ class PShellCommand(PCommand):
             for var in sorted(self.object_help.keys()):
                 help += '\n  %-12s %s' % (var, self.object_help[var])
 
+        if self.use_script:
+            help += '\n\nAll objects from %s are available.' % self.use_script
+
         if shell is None and not self.options.disable_ipython:
             shell = self.make_ipython_v0_11_shell()
             if shell is None:
diff --git a/pyramid/tests/pshellapp/__init__.py b/pyramid/tests/pshellapp/__init__.py
new file mode 100644
index 000000000..dd0a3e2f9
--- /dev/null
+++ b/pyramid/tests/pshellapp/__init__.py
@@ -0,0 +1,7 @@
+__all__ = ['a', 'root', 'm']
+
+a = 1
+b = 2
+
+root = 'root override'
+m = 'model override'
diff --git a/pyramid/tests/test_paster.py b/pyramid/tests/test_paster.py
index 36c3a51be..45e36f129 100644
--- a/pyramid/tests/test_paster.py
+++ b/pyramid/tests/test_paster.py
@@ -156,6 +156,46 @@ class TestPShellCommand(unittest.TestCase):
         self.assertTrue(self.bootstrap.closer.called)
         self.assertTrue(shell.help)
 
+    def test_command_loads_use_script(self):
+        command = self._makeOne()
+        self.config_factory.items = [('import', 'pyramid.tests.pshellapp')]
+        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 override',
+        })
+        self.assertTrue(self.bootstrap.closer.called)
+        self.assertTrue(shell.help)
+
+    def test_command_loads_use_script_check_order(self):
+        command = self._makeOne()
+        model = Dummy()
+        self.config_factory.items = [('import', 'pyramid.tests.pshellapp'),
+                                     ('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_custom_section_override(self):
         command = self._makeOne()
         dummy = Dummy()