diff options
Diffstat (limited to 'factory')
-rw-r--r-- | factory/__init__.py | 2 | ||||
-rw-r--r-- | factory/django.py | 60 | ||||
-rw-r--r-- | factory/helpers.py | 5 |
3 files changed, 67 insertions, 0 deletions
diff --git a/factory/__init__.py b/factory/__init__.py index b4e63be..251306a 100644 --- a/factory/__init__.py +++ b/factory/__init__.py @@ -79,5 +79,7 @@ from .helpers import ( lazy_attribute_sequence, container_attribute, post_generation, + + prevent_signals, ) diff --git a/factory/django.py b/factory/django.py index fee8e52..b502923 100644 --- a/factory/django.py +++ b/factory/django.py @@ -25,6 +25,9 @@ from __future__ import absolute_import from __future__ import unicode_literals import os +import types +import logging +import functools """factory_boy extensions for use with the Django framework.""" @@ -39,6 +42,9 @@ from . import base from . import declarations from .compat import BytesIO, is_string +logger = logging.getLogger('factory.generate') + + def require_django(): """Simple helper to ensure Django is available.""" @@ -214,3 +220,57 @@ class ImageField(FileField): thumb.save(thumb_io, format=image_format) return thumb_io.getvalue() + +class PreventSignals(object): + """Temporarily disables and then restores any django signals. + + Args: + *signals (django.dispatch.dispatcher.Signal): any django signals + + Examples: + with prevent_signals(pre_init): + user = UserFactory.build() + ... + + @prevent_signals(pre_save, post_save) + class UserFactory(factory.Factory): + ... + + @prevent_signals(post_save) + def generate_users(): + UserFactory.create_batch(10) + """ + + def __init__(self, *signals): + self.signals = signals + self.paused = {} + + def __enter__(self): + for signal in self.signals: + logger.debug('PreventSignals: Disabling signal handlers %r', + signal.receivers) + + self.paused[signal] = signal.receivers + signal.receivers = [] + + def __exit__(self, exc_type, exc_value, traceback): + for signal, receivers in self.paused.items(): + logger.debug('PreventSignals: Restoring signal handlers %r', + receivers) + + signal.receivers = receivers + self.paused = {} + + def __call__(self, func): + if isinstance(func, types.FunctionType): + @functools.wraps(func) + def wrapper(*args, **kwargs): + with self: + return func(*args, **kwargs) + return wrapper + + generate_method = getattr(func, '_generate', None) + if generate_method: + func._generate = classmethod(self(generate_method.__func__)) + + return func
\ No newline at end of file diff --git a/factory/helpers.py b/factory/helpers.py index 37b41bf..719d76d 100644 --- a/factory/helpers.py +++ b/factory/helpers.py @@ -28,6 +28,7 @@ import logging from . import base from . import declarations +from . import django @contextlib.contextmanager @@ -139,3 +140,7 @@ def container_attribute(func): def post_generation(fun): return declarations.PostGeneration(fun) + + +def prevent_signals(*signals): + return django.PreventSignals(*signals)
\ No newline at end of file |