Use extracted argument in PostGenerationMethodCall.

This changeset makes it possible possible to override the default method
arguments (or "method_args") passed in when instantiating
PostGenerationMethodCall. Now the user can override the default
arguments to the method called during post-generation when instantiating
a factory.

For example, using this UserFactory,

    class UserFactory(factory.Factory):
        FACTORY_FOR = User

        username = factory.Sequence(lambda n: 'user{0}'.format(n))
        password = factory.PostGenerationMethodCall(
                'set_password', None, 'defaultpassword')

by default, the user will have a password set to 'defaultpassword', but
this can be overridden by passing in a new password as a keyword
argument:

    >>> u = UserFactory()
    >>> u.check_password('defaultpassword')
    True
    >>> other_u = UserFactory(password='different')
    >>> other_u.check_password('defaultpassword')
    False
    >>> other_u.check_password('different')
    True

This changeset introduces a testing dependency on the Mock package
http://pypi.python.org/pypi/mock. While this is a third-party dependency
in Python 2, it is part of the Python 3 standard library, as unit.mock,
and so a reasonable dependency to satisfy.

Signed-off-by: Raphaël Barrois <raphael.barrois@polytechnique.org>

1 files changed, 9 insertions, 1 deletions

diff --git a/factory/declarations.py b/factory/declarations.py
index 366c2c8..1f1d2af 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -21,6 +21,7 @@
 # THE SOFTWARE.
 
 
+import collections
 import itertools
 import warnings
 
@@ -498,10 +499,17 @@ class PostGenerationMethodCall(PostGenerationDeclaration):
         self.method_kwargs = kwargs
 
     def call(self, obj, create, extracted=None, **kwargs):
+        if extracted is not None:
+            passed_args = extracted
+            if isinstance(passed_args, basestring) or (
+                    not isinstance(passed_args, collections.Iterable)):
+                passed_args = (passed_args,)
+        else:
+            passed_args = self.method_args
         passed_kwargs = dict(self.method_kwargs)
         passed_kwargs.update(kwargs)
         method = getattr(obj, self.method_name)
-        method(*self.method_args, **passed_kwargs)
+        method(*passed_args, **passed_kwargs)
 
 
 # Decorators... in case lambdas don't cut it