diff options
Diffstat (limited to 'factory')
| -rw-r--r-- | factory/__init__.py | 1 | ||||
| -rw-r--r-- | factory/base.py | 44 | ||||
| -rw-r--r-- | factory/declarations.py | 6 |
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): |