summaryrefslogtreecommitdiff
path: root/factory
diff options
context:
space:
mode:
Diffstat (limited to 'factory')
-rw-r--r--factory/__init__.py1
-rw-r--r--factory/base.py44
-rw-r--r--factory/declarations.py6
3 files changed, 39 insertions, 12 deletions
diff --git a/factory/__init__.py b/factory/__init__.py
index dd91343..1bf8968 100644
--- a/factory/__init__.py
+++ b/factory/__init__.py
@@ -67,5 +67,6 @@ from declarations import (
sequence,
lazy_attribute_sequence,
container_attribute,
+ post_declaration,
)
diff --git a/factory/base.py b/factory/base.py
index 98029e3..c1dbd98 100644
--- a/factory/base.py
+++ b/factory/base.py
@@ -44,6 +44,7 @@ FACTORY_CLASS_DECLARATION = 'FACTORY_FOR'
# Factory class attributes
CLASS_ATTRIBUTE_DECLARATIONS = '_declarations'
+CLASS_ATTRIBUTE_POSTGEN_DECLARATIONS = '_postgen_declarations'
CLASS_ATTRIBUTE_ASSOCIATED_CLASS = '_associated_class'
@@ -97,18 +98,24 @@ class BaseFactoryMetaClass(type):
return super(BaseFactoryMetaClass, cls).__new__(cls, class_name, bases, attrs)
declarations = containers.DeclarationDict()
+ postgen_declarations = containers.PostGenerationDeclarationDict()
# Add parent declarations in reverse order.
for base in reversed(parent_factories):
+ # Import parent PostGenerationDeclaration
+ postgen_declarations.update_with_public(
+ getattr(base, CLASS_ATTRIBUTE_POSTGEN_DECLARATIONS, {}))
# Import all 'public' attributes (avoid those starting with _)
declarations.update_with_public(getattr(base, CLASS_ATTRIBUTE_DECLARATIONS, {}))
- # Import attributes from the class definition, storing protected/private
- # attributes in 'non_factory_attrs'.
- non_factory_attrs = declarations.update_with_public(attrs)
+ # Import attributes from the class definition
+ non_postgen_attrs = postgen_declarations.update_with_public(attrs)
+ # Store protected/private attributes in 'non_factory_attrs'.
+ non_factory_attrs = declarations.update_with_public(non_postgen_attrs)
# Store the DeclarationDict in the attributes of the newly created class
non_factory_attrs[CLASS_ATTRIBUTE_DECLARATIONS] = declarations
+ non_factory_attrs[CLASS_ATTRIBUTE_POSTGEN_DECLARATIONS] = postgen_declarations
# Add extra args if provided.
if extra_attrs:
@@ -521,20 +528,37 @@ class Factory(BaseFactory):
return cls.get_building_function()(getattr(cls, CLASS_ATTRIBUTE_ASSOCIATED_CLASS), **kwargs)
@classmethod
- def _build(cls, **kwargs):
- return cls._prepare(create=False, **kwargs)
+ def _generate(cls, create, attrs):
+ """generate the object.
- @classmethod
- def _create(cls, **kwargs):
- return cls._prepare(create=True, **kwargs)
+ Args:
+ create (bool): whether to 'build' or 'create' the object
+ attrs (dict): attributes to use for generating the object
+ """
+ # Extract declarations used for post-generation
+ postgen_declarations = getattr(cls, CLASS_ATTRIBUTE_POSTGEN_DECLARATIONS)
+ postgen_attributes = {}
+ for name, decl in sorted(postgen_declarations.items()):
+ postgen_attributes[name] = decl.extract(name, attrs)
+
+ # Generate the object
+ obj = cls._prepare(create, **attrs)
+
+ # Handle post-generation attributes
+ for name, decl in sorted(postgen_declarations.items()):
+ extracted, extracted_kwargs = postgen_attributes[name]
+ decl.call(obj, create, extracted, **extracted_kwargs)
+ return obj
@classmethod
def build(cls, **kwargs):
- return cls._build(**cls.attributes(create=False, extra=kwargs))
+ attrs = cls.attributes(create=False, extra=kwargs)
+ return cls._generate(False, attrs)
@classmethod
def create(cls, **kwargs):
- return cls._create(**cls.attributes(create=True, extra=kwargs))
+ attrs = cls.attributes(create=True, extra=kwargs)
+ return cls._generate(True, attrs)
class DjangoModelFactory(Factory):
diff --git a/factory/declarations.py b/factory/declarations.py
index d9d560a..ddc6c78 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -23,6 +23,8 @@
import itertools
+from factory import utils
+
class OrderedDeclaration(object):
"""A factory declaration.
@@ -275,12 +277,12 @@ class PostGenerationDeclaration(object):
(object, dict): a tuple containing the attribute at 'name' (if
provided) and a dict of extracted attributes
"""
- extracted = attrs.get(name)
+ extracted = attrs.pop(name, None)
if self.extract_prefix:
extract_prefix = self.extract_prefix
else:
extract_prefix = name
- kwargs = utils.extract_dict(extract_prefix)
+ kwargs = utils.extract_dict(extract_prefix, attrs)
return extracted, kwargs
def call(self, obj, create, extracted=None, **kwargs):