diff options
Diffstat (limited to 'factory')
| -rw-r--r-- | factory/base.py | 18 | ||||
| -rw-r--r-- | factory/django.py | 19 |
2 files changed, 20 insertions, 17 deletions
diff --git a/factory/base.py b/factory/base.py index 8183649..3c6571c 100644 --- a/factory/base.py +++ b/factory/base.py @@ -175,11 +175,13 @@ class FactoryMetaClass(type): is_abstract = attrs.pop('ABSTRACT_FACTORY', False) base = parent_factories[0] - inherited_associated_class = getattr(base, - CLASS_ATTRIBUTE_ASSOCIATED_CLASS, None) + inherited_associated_class = base._get_target_class() associated_class = mcs._discover_associated_class(class_name, attrs, inherited_associated_class) + # Invoke 'lazy-loading' hooks. + associated_class = base._load_target_class(associated_class) + if associated_class is None: is_abstract = True @@ -379,13 +381,19 @@ class BaseFactory(object): return kwargs @classmethod - def _load_target_class(cls): + def _load_target_class(cls, class_definition): """Extension point for loading target classes. This can be overridden in framework-specific subclasses to hook into existing model repositories, for instance. """ - return getattr(cls, CLASS_ATTRIBUTE_ASSOCIATED_CLASS) + return class_definition + + @classmethod + def _get_target_class(cls): + """Retrieve the actual, associated target class.""" + definition = getattr(cls, CLASS_ATTRIBUTE_ASSOCIATED_CLASS, None) + return cls._load_target_class(definition) @classmethod def _prepare(cls, create, **kwargs): @@ -395,7 +403,7 @@ class BaseFactory(object): create: bool, whether to create or to build the object **kwargs: arguments to pass to the creation function """ - target_class = cls._load_target_class() + target_class = cls._get_target_class() kwargs = cls._adjust_kwargs(**kwargs) # Remove 'hidden' arguments. diff --git a/factory/django.py b/factory/django.py index 016586d..fee8e52 100644 --- a/factory/django.py +++ b/factory/django.py @@ -58,20 +58,15 @@ class DjangoModelFactory(base.Factory): ABSTRACT_FACTORY = True # Optional, but explicit. FACTORY_DJANGO_GET_OR_CREATE = () - _associated_model = None - @classmethod - def _load_target_class(cls): - associated_class = super(DjangoModelFactory, cls)._load_target_class() + def _load_target_class(cls, definition): - if is_string(associated_class) and '.' in associated_class: - app, model = associated_class.split('.', 1) - if cls._associated_model is None: - from django.db.models import loading as django_loading - cls._associated_model = django_loading.get_model(app, model) - return cls._associated_model + if is_string(definition) and '.' in definition: + app, model = definition.split('.', 1) + from django.db.models import loading as django_loading + return django_loading.get_model(app, model) - return associated_class + return definition @classmethod def _get_manager(cls, target_class): @@ -84,7 +79,7 @@ class DjangoModelFactory(base.Factory): def _setup_next_sequence(cls): """Compute the next available PK, based on the 'pk' database field.""" - model = cls._load_target_class() # pylint: disable=E1101 + model = cls._get_target_class() # pylint: disable=E1101 manager = cls._get_manager(model) try: |