# -*- coding: utf-8 -*-
# Copyright (c) 2010 Mark Sandstrom
# Copyright (c) 2011 Raphaël Barrois
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
import re
import sys
import warnings
from factory import containers
# Strategies
BUILD_STRATEGY = 'build'
CREATE_STRATEGY = 'create'
STUB_STRATEGY = 'stub'
# Creation functions. Use Factory.set_creation_function() to set a creation function appropriate for your ORM.
DJANGO_CREATION = lambda class_to_create, **kwargs: class_to_create.objects.create(**kwargs)
# Building functions. Use Factory.set_building_function() to set a building functions appropriate for your ORM.
NAIVE_BUILD = lambda class_to_build, **kwargs: class_to_build(**kwargs)
MOGO_BUILD = lambda class_to_build, **kwargs: class_to_build.new(**kwargs)
# Special declarations
FACTORY_CLASS_DECLARATION = 'FACTORY_FOR'
# Factory class attributes
CLASS_ATTRIBUTE_DECLARATIONS = '_declarations'
CLASS_ATTRIBUTE_POSTGEN_DECLARATIONS = '_postgen_declarations'
CLASS_ATTRIBUTE_ASSOCIATED_CLASS = '_associated_class'
# Factory metaclasses
[docs]def get_factory_bases(bases):
"""Retrieve all BaseFactoryMetaClass-derived bases from a list."""
return [b for b in bases if isinstance(b, BaseFactoryMetaClass)]
[docs]class BaseFactory(object):
"""Factory base support for sequences, attributes and stubs."""
class UnknownStrategy(RuntimeError):
pass
class UnsupportedStrategy(RuntimeError):
pass
def __new__(cls, *args, **kwargs):
"""Would be called if trying to instantiate the class."""
raise RuntimeError('You cannot instantiate BaseFactory')
# ID to use for the next 'declarations.Sequence' attribute.
_next_sequence = None
# Base factory, if this class was inherited from another factory. This is
# used for sharing the _next_sequence counter among factories for the same
# class.
_base_factory = None
@classmethod
def _setup_next_sequence(cls):
"""Set up an initial sequence value for Sequence attributes.
Returns:
int: the first available ID to use for instances of this factory.
"""
return 0
@classmethod
def _generate_next_sequence(cls):
"""Retrieve a new sequence ID.
This will call, in order:
- _generate_next_sequence from the base factory, if provided
- _setup_next_sequence, if this is the 'toplevel' factory and the
sequence counter wasn't initialized yet; then increase it.
"""
# Rely upon our parents
if cls._base_factory:
return cls._base_factory._generate_next_sequence()
# Make sure _next_sequence is initialized
if cls._next_sequence is None:
cls._next_sequence = cls._setup_next_sequence()
# Pick current value, then increase class counter for the next call.
next_sequence = cls._next_sequence
cls._next_sequence += 1
return next_sequence
@classmethod
[docs] def attributes(cls, create=False, extra=None):
"""Build a dict of attribute values, respecting declaration order.
The process is:
- Handle 'orderless' attributes, overriding defaults with provided
kwargs when applicable
- Handle ordered attributes, overriding them with provided kwargs when
applicable; the current list of computed attributes is available
to the currently processed object.
"""
return containers.AttributeBuilder(cls, extra).build(create)
@classmethod
[docs] def declarations(cls, extra_defs=None):
"""Retrieve a copy of the declared attributes.
Args:
extra_defs (dict): additional definitions to insert into the
retrieved DeclarationDict.
"""
return getattr(cls, CLASS_ATTRIBUTE_DECLARATIONS).copy(extra_defs)
@classmethod
[docs] def build(cls, **kwargs):
"""Build an instance of the associated class, with overriden attrs."""
raise cls.UnsupportedStrategy()
@classmethod
[docs] def build_batch(cls, size, **kwargs):
"""Build a batch of instances of the given class, with overriden attrs.
Args:
size (int): the number of instances to build
Returns:
object list: the built instances
"""
return [cls.build(**kwargs) for _ in xrange(size)]
@classmethod
[docs] def create(cls, **kwargs):
"""Create an instance of the associated class, with overriden attrs."""
raise cls.UnsupportedStrategy()
@classmethod
[docs] def create_batch(cls, size, **kwargs):
"""Create a batch of instances of the given class, with overriden attrs.
Args:
size (int): the number of instances to create
Returns:
object list: the created instances
"""
return [cls.create(**kwargs) for _ in xrange(size)]
@classmethod
[docs] def stub(cls, **kwargs):
"""Retrieve a stub of the associated class, with overriden attrs.
This will return an object whose attributes are those defined in this
factory's declarations or in the extra kwargs.
"""
stub_object = containers.StubObject()
for name, value in cls.attributes(create=False, extra=kwargs).iteritems():
setattr(stub_object, name, value)
return stub_object
@classmethod
[docs] def stub_batch(cls, size, **kwargs):
"""Stub a batch of instances of the given class, with overriden attrs.
Args:
size (int): the number of instances to stub
Returns:
object list: the stubbed instances
"""
return [cls.stub(**kwargs) for _ in xrange(size)]
@classmethod
[docs] def generate(cls, strategy, **kwargs):
"""Generate a new instance.
The instance will be created with the given strategy (one of
BUILD_STRATEGY, CREATE_STRATEGY, STUB_STRATEGY).
Args:
strategy (str): the strategy to use for generating the instance.
Returns:
object: the generated instance
"""
assert strategy in (STUB_STRATEGY, BUILD_STRATEGY, CREATE_STRATEGY)
action = getattr(cls, strategy)
return action(**kwargs)
@classmethod
[docs] def generate_batch(cls, strategy, size, **kwargs):
"""Generate a batch of instances.
The instances will be created with the given strategy (one of
BUILD_STRATEGY, CREATE_STRATEGY, STUB_STRATEGY).
Args:
strategy (str): the strategy to use for generating the instance.
size (int): the number of instances to generate
Returns:
object list: the generated instances
"""
assert strategy in (STUB_STRATEGY, BUILD_STRATEGY, CREATE_STRATEGY)
batch_action = getattr(cls, '%s_batch' % strategy)
return batch_action(size, **kwargs)
@classmethod
[docs] def simple_generate(cls, create, **kwargs):
"""Generate a new instance.
The instance will be either 'built' or 'created'.
Args:
create (bool): whether to 'build' or 'create' the instance.
Returns:
object: the generated instance
"""
strategy = CREATE_STRATEGY if create else BUILD_STRATEGY
return cls.generate(strategy, **kwargs)
@classmethod
[docs] def simple_generate_batch(cls, create, size, **kwargs):
"""Generate a batch of instances.
These instances will be either 'built' or 'created'.
Args:
size (int): the number of instances to generate
create (bool): whether to 'build' or 'create' the instances.
Returns:
object list: the generated instances
"""
strategy = CREATE_STRATEGY if create else BUILD_STRATEGY
return cls.generate_batch(strategy, size, **kwargs)
class StubFactory(BaseFactory):
__metaclass__ = BaseFactoryMetaClass
default_strategy = STUB_STRATEGY
[docs]class Factory(BaseFactory):
"""Factory base with build and create support.
This class has the ability to support multiple ORMs by using custom creation
functions.
"""
__metaclass__ = FactoryMetaClass
default_strategy = CREATE_STRATEGY
class AssociatedClassError(RuntimeError):
pass
# Customizing 'create' strategy, using a tuple to keep the creation function
# from turning it into an instance method.
_creation_function = (DJANGO_CREATION,)
@classmethod
[docs] def set_creation_function(cls, creation_function):
"""Set the creation function for this class.
Args:
creation_function (function): the new creation function. That
function should take one non-keyword argument, the 'class' for
which an instance will be created. The value of the various
fields are passed as keyword arguments.
"""
cls._creation_function = (creation_function,)
@classmethod
[docs] def get_creation_function(cls):
"""Retrieve the creation function for this class.
Returns:
function: A function that takes one parameter, the class for which
an instance will be created, and keyword arguments for the value
of the fields of the instance.
"""
return cls._creation_function[0]
# Customizing 'build' strategy, using a tuple to keep the creation function
# from turning it into an instance method.
_building_function = (NAIVE_BUILD,)
@classmethod
[docs] def set_building_function(cls, building_function):
"""Set the building function for this class.
Args:
building_function (function): the new building function. That
function should take one non-keyword argument, the 'class' for
which an instance will be built. The value of the various
fields are passed as keyword arguments.
"""
cls._building_function = (building_function,)
@classmethod
[docs] def get_building_function(cls):
"""Retrieve the building function for this class.
Returns:
function: A function that takes one parameter, the class for which
an instance will be created, and keyword arguments for the value
of the fields of the instance.
"""
return cls._building_function[0]
@classmethod
def _prepare(cls, create, **kwargs):
"""Prepare an object for this factory.
Args:
create: bool, whether to create or to build the object
**kwargs: arguments to pass to the creation function
"""
if create:
return cls.get_creation_function()(getattr(cls, CLASS_ATTRIBUTE_ASSOCIATED_CLASS), **kwargs)
else:
return cls.get_building_function()(getattr(cls, CLASS_ATTRIBUTE_ASSOCIATED_CLASS), **kwargs)
@classmethod
def _generate(cls, create, attrs):
"""generate the object.
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):
attrs = cls.attributes(create=False, extra=kwargs)
return cls._generate(False, attrs)
@classmethod
def create(cls, **kwargs):
attrs = cls.attributes(create=True, extra=kwargs)
return cls._generate(True, attrs)
[docs]class DjangoModelFactory(Factory):
"""Factory for Django models.
This makes sure that the 'sequence' field of created objects is an unused id.
Possible improvement: define a new 'attribute' type, AutoField, which would
handle those for non-numerical primary keys.
"""
ABSTRACT_FACTORY = True
@classmethod
def _setup_next_sequence(cls):
"""Compute the next available ID, based on the 'id' database field."""
try:
return 1 + cls._associated_class._default_manager.values_list('id', flat=True
).order_by('-id')[0]
except IndexError:
return 1
[docs]def make_factory(klass, **kwargs):
"""Create a new, simple factory for the given class."""
factory_name = '%sFactory' % klass.__name__
kwargs[FACTORY_CLASS_DECLARATION] = klass
factory_class = type(Factory).__new__(type(Factory), factory_name, (Factory,), kwargs)
factory_class.__name__ = '%sFactory' % klass.__name__
factory_class.__doc__ = 'Auto-generated factory for class %s' % klass
return factory_class
[docs]def build(klass, **kwargs):
"""Create a factory for the given class, and build an instance."""
return make_factory(klass, **kwargs).build()
[docs]def build_batch(klass, size, **kwargs):
"""Create a factory for the given class, and build a batch of instances."""
return make_factory(klass, **kwargs).build_batch(size)
[docs]def create(klass, **kwargs):
"""Create a factory for the given class, and create an instance."""
return make_factory(klass, **kwargs).create()
[docs]def create_batch(klass, size, **kwargs):
"""Create a factory for the given class, and create a batch of instances."""
return make_factory(klass, **kwargs).create_batch(size)
[docs]def stub(klass, **kwargs):
"""Create a factory for the given class, and stub an instance."""
return make_factory(klass, **kwargs).stub()
[docs]def stub_batch(klass, size, **kwargs):
"""Create a factory for the given class, and stub a batch of instances."""
return make_factory(klass, **kwargs).stub_batch(size)
[docs]def generate(klass, strategy, **kwargs):
"""Create a factory for the given class, and generate an instance."""
return make_factory(klass, **kwargs).generate(strategy)
[docs]def generate_batch(klass, strategy, size, **kwargs):
"""Create a factory for the given class, and generate instances."""
return make_factory(klass, **kwargs).generate_batch(strategy, size)
[docs]def simple_generate(klass, create, **kwargs):
"""Create a factory for the given class, and simple_generate an instance."""
return make_factory(klass, **kwargs).simple_generate(create)
[docs]def simple_generate_batch(klass, create, size, **kwargs):
"""Create a factory for the given class, and simple_generate instances."""
return make_factory(klass, **kwargs).simple_generate_batch(create, size)
[docs]def use_strategy(new_strategy):
"""Force the use of a different strategy.
This is an alternative to setting default_strategy in the class definition.
"""
def wrapped_class(klass):
klass.default_strategy = new_strategy
return klass
return wrapped_class