From b36dd80c4400dbaacf9809d9bc1c9147631382f5 Mon Sep 17 00:00:00 2001
From: Michael Merickel <michael@merickel.org>
Date: Thu, 1 Sep 2016 01:07:49 -0500
Subject: [PATCH] allow prepare/bootstrap to be used as a context manager

closes #1822
---
 pyramid/paster.py               | 14 +++++++++++++
 pyramid/scripting.py            | 35 ++++++++++++++++++++++++++++-----
 pyramid/tests/test_scripting.py | 21 ++++++++++++++++++++
 3 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/pyramid/paster.py b/pyramid/paster.py
index 3916be8f09..1b7afb5dcb 100644
--- a/pyramid/paster.py
+++ b/pyramid/paster.py
@@ -129,8 +129,22 @@ def bootstrap(config_uri, request=None, options=None):
     {'http_port': 8080} and then use %(http_port)s in the
     config file.
 
+    This function may be used as a context manager to call the ``closer``
+    automatically:
+
+    .. code-block:: python
+
+       with bootstrap('development.ini') as env:
+           request = env['request']
+           # ...
+
     See :ref:`writing_a_script` for more information about how to use this
     function.
+
+    .. versionchanged:: 1.8
+
+       Added the ability to use the return value as a context manager.
+
     """
     app = get_app(config_uri, options=options)
     env = prepare(request)
diff --git a/pyramid/scripting.py b/pyramid/scripting.py
index d9587338ff..7607d3ea3f 100644
--- a/pyramid/scripting.py
+++ b/pyramid/scripting.py
@@ -56,12 +56,25 @@ def prepare(request=None, registry=None):
     ``root`` returned is the application's root resource object.  The
     ``closer`` returned is a callable (accepting no arguments) that
     should be called when your scripting application is finished
-    using the root.  ``registry`` is the registry object passed or
-    the last registry loaded into
-    :attr:`pyramid.config.global_registries` if no registry is passed.
+    using the root.  ``registry`` is the resolved registry object.
     ``request`` is the request object passed or the constructed request
     if no request is passed.  ``root_factory`` is the root factory used
     to construct the root.
+
+    This function may be used as a context manager to call the ``closer``
+    automatically:
+
+    .. code-block:: python
+
+       registry = config.registry
+       with prepare(registry) as env:
+           request = env['request']
+           # ...
+
+    .. versionchanged:: 1.8
+
+       Added the ability to use the return value as a context manager.
+
     """
     if registry is None:
         registry = getattr(request, 'registry', global_registries.last)
@@ -85,8 +98,20 @@ def closer():
     root = root_factory(request)
     if getattr(request, 'context', None) is None:
         request.context = root
-    return {'root':root, 'closer':closer, 'registry':registry,
-            'request':request, 'root_factory':root_factory}
+    return AppEnvironment(
+        root=root,
+        closer=closer,
+        registry=registry,
+        request=request,
+        root_factory=root_factory,
+    )
+
+class AppEnvironment(dict):
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self['closer']()
 
 def _make_request(path, registry=None):
     """ Return a :meth:`pyramid.request.Request` object anchored at a
diff --git a/pyramid/tests/test_scripting.py b/pyramid/tests/test_scripting.py
index 1e952062ba..00f738e026 100644
--- a/pyramid/tests/test_scripting.py
+++ b/pyramid/tests/test_scripting.py
@@ -134,6 +134,27 @@ def test_it_with_extensions(self):
         root, closer = info['root'], info['closer']
         closer()
 
+    def test_it_is_a_context_manager(self):
+        request = DummyRequest({})
+        registry = request.registry = self._makeRegistry()
+        closer_called = [False]
+        with self._callFUT(request=request) as info:
+            root, request = info['root'], info['request']
+            pushed = self.manager.get()
+            self.assertEqual(pushed['request'], request)
+            self.assertEqual(pushed['registry'], registry)
+            self.assertEqual(pushed['request'].registry, registry)
+            self.assertEqual(root.a, (request,))
+            orig_closer = info['closer']
+            def closer():
+                orig_closer()
+                closer_called[0] = True
+            info['closer'] = closer
+        self.assertTrue(closer_called[0])
+        self.assertEqual(self.default, self.manager.get())
+        self.assertEqual(request.context, root)
+        self.assertEqual(request.registry, registry)
+
 class Test__make_request(unittest.TestCase):
     def _callFUT(self, path='/', registry=None):
         from pyramid.scripting import _make_request