Added support for object ensuring

4 files changed, 155 insertions, 7 deletions

diff --git a/click/core.py b/click/core.py
index df97726..e0eb78d 100644
--- a/click/core.py
+++ b/click/core.py
@@ -133,6 +133,15 @@ class Context(object):
                 return node.obj
             node = node.parent
 
+    def ensure_object(self, object_type):
+        """Like :meth:`find_object` but sets the innermost object to a
+        new instance of `object_type` if it does not exist.
+        """
+        rv = self.find_object(object_type)
+        if rv is None:
+            self.obj = rv = object_type()
+        return rv
+
     def fail(self, message):
         """Aborts the execution of the program with a specific error
         message.
diff --git a/click/decorators.py b/click/decorators.py
index 08cecc8..aeeb749 100644
--- a/click/decorators.py
+++ b/click/decorators.py
@@ -25,7 +25,7 @@ def pass_obj(f):
     return update_wrapper(new_func, f)
 
 
-def make_pass_decorator(object_type):
+def make_pass_decorator(object_type, ensure=False):
     """Given an object type this creates a decorator that will work
     similar to :func:`pass_obj` but instead of passing the object of the
     current context, it will find the innermost context of type
@@ -42,12 +42,19 @@ def make_pass_decorator(object_type):
                 return ctx.invoke(f, obj, *args, **kwargs)
             return update_wrapper(new_func, f)
         return decorator
+
+    :param object_type: the type of the object to pass.
+    :param ensure: if set to `True`, a new object will be created and
+                   remembered on the context if it's not there yet.
     """
     def decorator(f):
         @pass_context
         def new_func(*args, **kwargs):
             ctx = args[0]
-            obj = ctx.find_object(object_type)
+            if ensure:
+                obj = ctx.ensure_object(object_type)
+            else:
+                obj = ctx.find_object(object_type)
             if obj is None:
                 raise RuntimeError('Managed to invoke callback without a '
                                    'context object of type %r existing'
diff --git a/docs/complex.rst b/docs/complex.rst
index a04a6e2..ebf1efb 100644
--- a/docs/complex.rst
+++ b/docs/complex.rst
@@ -82,8 +82,8 @@ state of our tool:
 
 
     class Repo(object):
-        def __init__(self, home, debug):
-            self.home = os.path.abspath(home)
+        def __init__(self, home=None, debug=False):
+            self.home = os.path.abspath(home or '.')
             self.debug = debug
 
 
@@ -159,9 +159,10 @@ created our repo.  Because of that we can start a search for the last
 level where the object stored on the context was a repo.
 
 Built-in support for this is provided by the :func:`make_pass_decorator`
-factory which will create decorators for us that find objects.  So in our
-case we know that we want to find the closest ``Repo`` object.  So let's
-make a decorator for this:
+factory which will create decorators for us that find objects (it
+internally calls into :meth:`Context.find_object`).  So in our case we
+know that we want to find the closest ``Repo`` object.  So let's make a
+decorator for this:
 
 .. click:example::
 
@@ -178,3 +179,41 @@ repo instead of something else:
     @pass_repo
     def clone(repo, src, dest):
         pass
+
+Ensuring Objects
+````````````````
+
+The above example only works if there was an outer command that created a
+``Repo`` object and stored it on the context.  For some more advanced use
+cases this might be a problem for you.  The default behavior of
+:func:`make_pass_decorator` is to call into :meth:`Context.find_object`
+which will find the object.  If it can't find the object it will raise an
+error.  The alternative behavior is to use :meth:`Context.ensure_object`
+which will find the object, or if it cannot find it, will create one and
+store it on the innermost context.  This behavior can also be enabled for
+:func:`make_pass_decorator` by passing ``ensure=True``:
+
+.. click:example::
+
+    pass_repo = click.make_pass_decorator(Repo, ensure=True)
+
+In this case the innermost context gets such an object created if it's
+missing.  This might replace objects being placed there earlier.  In this
+case the command stays executable, even if the outer command does not run.
+For this to work the object type needs to have a constructor that accepts
+no arguments.
+
+As such it runs standalone:
+
+.. click:example::
+
+    @click.command()
+    @pass_repo
+    def cp(repo):
+        click.echo(repo)
+
+As you can see:
+
+.. click:run::
+
+    invoke(cp, [])
diff --git a/tests/test_context.py b/tests/test_context.py
new file mode 100644
index 0000000..92ef447
--- /dev/null
+++ b/tests/test_context.py
@@ -0,0 +1,93 @@
+# -*- coding: utf-8 -*-
+import click
+
+
+def test_ensure_context_objects(runner):
+    class Foo(object):
+        def __init__(self):
+            self.title = 'default'
+
+    pass_foo = click.make_pass_decorator(Foo, ensure=True)
+
+    @click.group()
+    @pass_foo
+    def cli(foo):
+        pass
+
+    @cli.command()
+    @pass_foo
+    def test(foo):
+        click.echo(foo.title)
+
+    result = runner.invoke(cli, ['test'])
+    assert not result.exception
+    assert result.output == 'default\n'
+
+
+def test_get_context_objects(runner):
+    class Foo(object):
+        def __init__(self):
+            self.title = 'default'
+
+    pass_foo = click.make_pass_decorator(Foo, ensure=True)
+
+    @click.group()
+    @click.pass_context
+    def cli(ctx):
+        ctx.obj = Foo()
+        ctx.obj.title = 'test'
+
+    @cli.command()
+    @pass_foo
+    def test(foo):
+        click.echo(foo.title)
+
+    result = runner.invoke(cli, ['test'])
+    assert not result.exception
+    assert result.output == 'test\n'
+
+
+def test_get_context_objects_no_ensuring(runner):
+    class Foo(object):
+        def __init__(self):
+            self.title = 'default'
+
+    pass_foo = click.make_pass_decorator(Foo)
+
+    @click.group()
+    @click.pass_context
+    def cli(ctx):
+        ctx.obj = Foo()
+        ctx.obj.title = 'test'
+
+    @cli.command()
+    @pass_foo
+    def test(foo):
+        click.echo(foo.title)
+
+    result = runner.invoke(cli, ['test'])
+    assert not result.exception
+    assert result.output == 'test\n'
+
+
+def test_get_context_objects_missing(runner):
+    class Foo(object):
+        pass
+
+    pass_foo = click.make_pass_decorator(Foo)
+
+    @click.group()
+    @click.pass_context
+    def cli(ctx):
+        pass
+
+    @cli.command()
+    @pass_foo
+    def test(foo):
+        click.echo(foo.title)
+
+    result = runner.invoke(cli, ['test'])
+    assert result.exception is not None
+    assert isinstance(result.exception, RuntimeError)
+    assert "Managed to invoke callback without a context object " \
+        "of type 'Foo' existing" in str(result.exception)