Add Traits (Closes #251).

Based on a boolean flag, those will alter the definitions of the current
factory, taking precedence over pre-defined behavior but overridden by
callsite-level arguments.

6 files changed, 382 insertions, 55 deletions

diff --git a/README.rst b/README.rst
index a08d37f..0e1f000 100644
--- a/README.rst
+++ b/README.rst
@@ -292,6 +292,17 @@ The associated object's strategy will be used:
     True
 
 
+ORM Support
+"""""""""""
+
+factory_boy has specific support for a few ORMs, through specific ``factory.Factory`` subclasses:
+
+* Django, with ``factory.django.DjangoModelFactory``
+* Mogo, with ``factory.mogo.MogoFactory``
+* MongoEngine, with ``factory.mongoengine.MongoEngineFactory``
+* SQLAlchemy, with ``factory.alchemy.SQLAlchemyModelFactory``
+
+
 Debugging factory_boy
 """""""""""""""""""""
 
@@ -326,16 +337,6 @@ This will yield messages similar to those (artificial indentation):
     BaseFactory: Generating tests.test_using.TestModel2Factory(two=<tests.test_using.TestModel object at 0x1e15410>)
 
 
-ORM Support
-"""""""""""
-
-factory_boy has specific support for a few ORMs, through specific ``factory.Factory`` subclasses:
-
-* Django, with ``factory.django.DjangoModelFactory``
-* Mogo, with ``factory.mogo.MogoFactory``
-* MongoEngine, with ``factory.mongoengine.MongoEngineFactory``
-* SQLAlchemy, with ``factory.alchemy.SQLAlchemyModelFactory``
-
 Contributing
 ------------
 
diff --git a/docs/introduction.rst b/docs/introduction.rst
index 9a16c39..5b535c9 100644
--- a/docs/introduction.rst
+++ b/docs/introduction.rst
@@ -266,6 +266,61 @@ This is handled by the :data:`~factory.FactoryOptions.inline_args` attribute:
     <MyClass(1, 4, z=3)>
 
 
+Altering a factory's behaviour: parameters and traits
+-----------------------------------------------------
+
+Some classes are better described with a few, simple parameters, that aren't fields on the actual model.
+In that case, use a :attr:`~factory.Factory.Params` declaration:
+
+.. code-block:: python
+
+    class RentalFactory(factory.Factory):
+        class Meta:
+            model = Rental
+
+        begin = factory.fuzzy.FuzzyDate(start_date=datetime.date(2000, 1, 1))
+        end = factory.LazyAttribute(lambda o: o.begin + o.duration)
+
+        class Params:
+            duration = 12
+
+.. code-block:: pycon
+
+    >>> RentalFactory(duration=0)
+    <Rental: 2012-03-03 -> 2012-03-03>
+    >>> RentalFactory(duration=10)
+    <Rental: 2008-12-16 -> 2012-12-26>
+
+
+When many fields should be updated based on a flag, use :class:`Traits <factory.Trait>` instead:
+
+.. code-block:: python
+
+    class OrderFactory(factory.Factory):
+        status = 'pending'
+        shipped_by = None
+        shipped_on = None
+
+        class Meta:
+            model = Order
+
+        class Params:
+            shipped = factory.Trait(
+                status='shipped',
+                shipped_by=factory.SubFactory(EmployeeFactory),
+                shipped_on=factory.LazyFunction(datetime.date.today),
+            )
+
+A trait is toggled by a single boolean value:
+
+.. code-block:: pycon
+
+    >>> OrderFactory()
+    <Order: pending>
+    >>> OrderFactory(shipped=True)
+    <Order: shipped by John Doe on 2016-04-02>
+
+
 Strategies
 ----------
 
diff --git a/docs/reference.rst b/docs/reference.rst
index 8550f88..ad68faf 100644
--- a/docs/reference.rst
+++ b/docs/reference.rst
@@ -11,6 +11,9 @@ For internals and customization points, please refer to the :doc:`internals` sec
 The :class:`Factory` class
 --------------------------
 
+Meta options
+""""""""""""
+
 .. class:: FactoryOptions
 
     .. versionadded:: 2.4.0
@@ -135,11 +138,16 @@ The :class:`Factory` class
 
 
 
+Attributes and methods
+""""""""""""""""""""""
+
+
 .. class:: Factory
 
 
     **Class-level attributes:**
 
+    .. attribute:: Meta
     .. attribute:: _meta
 
         .. versionadded:: 2.4.0
@@ -147,6 +155,14 @@ The :class:`Factory` class
         The :class:`FactoryOptions` instance attached to a :class:`Factory` class is available
         as a :attr:`_meta` attribute.
 
+    .. attribute:: Params
+
+        .. versionadded:: 2.7.0
+
+        The extra parameters attached to a :class:`Factory` are declared through a :attr:`Params`
+        class.
+        See :ref:`the "Parameters" section <parameters>` for more information.
+
     .. attribute:: _options_class
 
         .. versionadded:: 2.4.0
@@ -353,6 +369,175 @@ The :class:`Factory` class
         factory in the chain.
 
 
+.. _parameters:
+
+Parameters
+""""""""""
+
+.. versionadded:: 2.7.0
+
+Some models have many fields that can be summarized by a few parameters; for instance,
+a train with many cars — each complete with serial number, manufacturer, ...;
+or an order that can be pending/shipped/received, with a few fields to describe each step.
+
+When building instances of such models, a couple of parameters can be enough to determine
+all other fields; this is handled by the :class:`~Factory.Params` section of a :class:`Factory` declaration.
+
+
+Simple parameters
+~~~~~~~~~~~~~~~~~
+
+Some factories only need little data:
+
+.. code-block:: python
+
+    class ConferenceFactory(factory.Factory):
+        class Meta:
+            model = Conference
+
+        class Params:
+            duration = 'short' # Or 'long'
+
+        start_date = factory.fuzzy.FuzzyDate()
+        end_date = factory.LazyAttribute(
+            lambda o: o.start_date + datetime.timedelta(days=2 if o.duration == 'short' else 7)
+        )
+        sprints_start = factory.LazyAttribute(
+            lambda o: o.end_date - datetime.timedelta(days=0 if o.duration == 'short' else 1)
+        )
+
+.. code-block:: pycon
+
+    >>> Conference(duration='short')
+    <Conference: DUTH 2015 (2015-11-05 - 2015-11-08, sprints 2015-11-08)>
+    >>> Conference(duration='long')
+    <Conference: DjangoConEU 2016 (2016-03-30 - 2016-04-03, sprints 2016-04-02)>
+
+
+Any simple parameter provided to the :class:`Factory.Params` section is available to the whole factory,
+but not passed to the final class (similar to the :attr:`~FactoryOptions.exclude` behavior).
+
+
+Traits
+~~~~~~
+
+.. class:: Trait(**kwargs)
+
+    .. OHAI VIM**
+
+    .. versionadded:: 2.7.0
+
+    A trait's parameters are the fields it sohuld alter when enabled.
+
+
+For more complex situations, it is helpful to override a few fields at once:
+
+.. code-block:: python
+
+    class OrderFactory(factory.Factory):
+        class Meta:
+            model = Order
+
+        state = 'pending'
+        shipped_on = None
+        shipped_by = None
+
+        class Params:
+            shipped = factory.Trait(
+                state='shipped',
+                shipped_on=datetime.date.today,
+                shipped_by=factory.SubFactory(EmployeeFactory),
+            )
+
+Such a :class:`Trait` is activated or disabled by a single boolean field:
+
+
+.. code-block:: pycon
+
+    >>> OrderFactory()
+    <Order: pending>
+    Order(state='pending')
+    >>> OrderFactory(shipped=True)
+    <Order: shipped by John Doe on 2016-04-02>
+
+
+A :class:`Trait` can be enabled/disabled by a :class:`Factory` subclass:
+
+.. code-block:: python
+
+    class ShippedOrderFactory(OrderFactory):
+        shipped = True
+
+
+Values set in a :class:`Trait` can be overridden by call-time values:
+
+.. code-block:: pycon
+
+    >>> OrderFactory(shipped=True, shipped_on=last_year)
+    <Order: shipped by John Doe on 2015-04-20>
+
+
+:class:`Traits <Trait>` can be chained:
+
+.. code-block:: python
+
+    class OrderFactory(factory.Factory):
+        class Meta:
+            model = Order
+
+        # Can be pending/shipping/received
+        state = 'pending'
+        shipped_on = None
+        shipped_by = None
+        received_on = None
+        received_by = None
+
+        class Params:
+            shipped = factory.Trait(
+                state='shipped',
+                shipped_on=datetime.date.today,
+                shipped_by=factory.SubFactory(EmployeeFactory),
+            )
+            received = factory.Trait(
+                shipped=True,
+                state='received',
+                shipped_on=datetime.date.today - datetime.timedelta(days=4),
+                received_on=datetime.date.today,
+                received_by=factory.SubFactory(CustomerFactory),
+            )
+
+.. code-block:: pycon
+
+    >>> OrderFactory(received=True)
+    <Order: shipped by John Doe on 2016-03-20, received by Joan Smith on 2016-04-02>
+
+
+
+A :class:`Trait` might be overridden in :class:`Factory` subclasses:
+
+.. code-block:: python
+
+    class LocalOrderFactory(OrderFactory):
+
+        class Params:
+            received = factory.Trait(
+                shipped=True,
+                state='received',
+                shipped_on=datetime.date.today - datetime.timedelta(days=1),
+                received_on=datetime.date.today,
+                received_by=factory.SubFactory(CustomerFactory),
+            )
+
+
+.. code-block:: pycon
+
+    >>> LocalOrderFactory(received=True)
+    <Order: shipped by John Doe on 2016-04-01, received by Joan Smith on 2016-04-02>
+
+
+.. note:: When overriding a :class:`Trait`, the whole declaration **MUST** be replaced.
+
+
 .. _strategies:
 
 Strategies
@@ -1299,51 +1484,6 @@ with the :class:`Dict` and :class:`List` attributes:
         argument, if another type (tuple, set, ...) is required.
 
 
-Parameters
-""""""""""
-
-Some models have many fields that can be summarized by a few parameters; for instance,
-a train with many cars — each complete with serial number, manufacturer, ...;
-or an order that can be pending/shipped/received, with a few fields to describe each step.
-
-When building instances of such models, a couple of parameters can be enough to determine
-all other fields; this is handled by the :class:`~Factory.Params` section of a :class:`Factory` declaration.
-
-
-Simple parameters
-~~~~~~~~~~~~~~~~~
-
-Some factories only need little data:
-
-.. code-block:: python
-
-    class ConferenceFactory(factory.Factory):
-        class Meta:
-            model = Conference
-
-        class Params:
-            duration = 'short' # Or 'long'
-
-        start_date = factory.fuzzy.FuzzyDate()
-        end_date = factory.LazyAttribute(
-            lambda o: o.start_date + datetime.timedelta(days=2 if o.duration == 'short' else 7)
-        )
-        sprints_start = factory.LazyAttribute(
-            lambda o: o.end_date - datetime.timedelta(days=0 if o.duration == 'short' else 1)
-        )
-
-.. code-block:: pycon
-
-    >>> Conference(duration='short')
-    <Conference: DUTH 2015 (2015-11-05 - 2015-11-08, sprints 2015-11-08)>
-    >>> Conference(duration='long')
-    <Conference: DjangoConEU 2016 (2016-03-30 - 2016-04-03, sprints 2016-04-02)>
-
-
-Any simple parameter provided to the :class:`Factory.Params` section is available to the whole factory,
-but not passed to the final class (similar to the :attr:`~FactoryOptions.exclude` behavior).
-
-
 
 Post-generation hooks
 """""""""""""""""""""
diff --git a/factory/__init__.py b/factory/__init__.py
index ccd71cd..ad9da80 100644
--- a/factory/__init__.py
+++ b/factory/__init__.py
@@ -52,6 +52,7 @@ from .declarations import (
     Sequence,
     LazyAttributeSequence,
     SelfAttribute,
+    Trait,
     ContainerAttribute,
     SubFactory,
     Dict,
diff --git a/factory/declarations.py b/factory/declarations.py
index ad1f72f..895f2ac 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -469,6 +469,22 @@ class ComplexParameter(object):
         return []
 
 
+class Trait(ComplexParameter):
+    """The simplest complex parameter, it enables a bunch of new declarations based on a boolean flag."""
+    def __init__(self, **overrides):
+        self.overrides = overrides
+
+    def compute(self, field_name, declarations):
+        if declarations.get(field_name):
+            return self.overrides
+        else:
+            return {}
+
+    def get_revdeps(self, parameters):
+        """This might alter fields it's injecting."""
+        return [param for param in parameters if param in self.overrides]
+
+
 # Post-generation
 # ===============
 
diff --git a/tests/test_using.py b/tests/test_using.py
index 67db3bc..eaeb8da 100644
--- a/tests/test_using.py
+++ b/tests/test_using.py
@@ -27,6 +27,7 @@ import sys
 import warnings
 
 import factory
+from factory import errors
 
 from .compat import is_python2, unittest
 from . import tools
@@ -1114,6 +1115,119 @@ class KwargAdjustTestCase(unittest.TestCase):
         self.assertEqual(42, obj.attributes)
 
 
+class TraitTestCase(unittest.TestCase):
+    def test_traits(self):
+        class TestObjectFactory(factory.Factory):
+            class Meta:
+                model = TestObject
+
+            class Params:
+                even = factory.Trait(two=True, four=True)
+                odd = factory.Trait(one=True, three=True, five=True)
+
+        obj1 = TestObjectFactory()
+        self.assertEqual(obj1.as_dict(),
+            dict(one=None, two=None, three=None, four=None, five=None))
+
+        obj2 = TestObjectFactory(even=True)
+        self.assertEqual(obj2.as_dict(),
+            dict(one=None, two=True, three=None, four=True, five=None))
+
+        obj3 = TestObjectFactory(odd=True)
+        self.assertEqual(obj3.as_dict(),
+            dict(one=True, two=None, three=True, four=None, five=True))
+
+        obj4 = TestObjectFactory(even=True, odd=True)
+        self.assertEqual(obj4.as_dict(),
+            dict(one=True, two=True, three=True, four=True, five=True))
+
+        obj5 = TestObjectFactory(odd=True, two=True)
+        self.assertEqual(obj5.as_dict(),
+            dict(one=True, two=True, three=True, four=None, five=True))
+
+    def test_traits_inheritance(self):
+        """A trait can be set in an inherited class."""
+        class TestObjectFactory(factory.Factory):
+            class Meta:
+                model = TestObject
+
+            class Params:
+                even = factory.Trait(two=True, four=True)
+                odd = factory.Trait(one=True, three=True, five=True)
+
+        class EvenObjectFactory(TestObjectFactory):
+            even = True
+
+        # Simple call
+        obj1 = EvenObjectFactory()
+        self.assertEqual(obj1.as_dict(),
+            dict(one=None, two=True, three=None, four=True, five=None))
+
+        # Force-disable it
+        obj2 = EvenObjectFactory(even=False)
+        self.assertEqual(obj2.as_dict(),
+            dict(one=None, two=None, three=None, four=None, five=None))
+
+    def test_traits_override(self):
+        """Override a trait in a subclass."""
+        class TestObjectFactory(factory.Factory):
+            class Meta:
+                model = TestObject
+
+            class Params:
+                even = factory.Trait(two=True, four=True)
+                odd = factory.Trait(one=True, three=True, five=True)
+
+        class WeirdMathFactory(TestObjectFactory):
+            class Params:
+                # Here, one is even.
+                even = factory.Trait(two=True, four=True, one=True)
+
+        obj = WeirdMathFactory(even=True)
+        self.assertEqual(obj.as_dict(),
+            dict(one=True, two=True, three=None, four=True, five=None))
+
+    def test_traits_chaining(self):
+        """Use a trait to enable other traits."""
+        class TestObjectFactory(factory.Factory):
+            class Meta:
+                model = TestObject
+
+            class Params:
+                even = factory.Trait(two=True, four=True)
+                odd = factory.Trait(one=True, three=True, five=True)
+                full = factory.Trait(even=True, odd=True)
+
+        # Setting "full" should enable all fields.
+        obj = TestObjectFactory(full=True)
+        self.assertEqual(obj.as_dict(),
+            dict(one=True, two=True, three=True, four=True, five=True))
+
+        # Does it break usual patterns?
+        obj1 = TestObjectFactory()
+        self.assertEqual(obj1.as_dict(),
+            dict(one=None, two=None, three=None, four=None, five=None))
+
+        obj2 = TestObjectFactory(even=True)
+        self.assertEqual(obj2.as_dict(),
+            dict(one=None, two=True, three=None, four=True, five=None))
+
+        obj3 = TestObjectFactory(odd=True)
+        self.assertEqual(obj3.as_dict(),
+            dict(one=True, two=None, three=True, four=None, five=True))
+
+    def test_prevent_cyclic_traits(self):
+
+        with self.assertRaises(errors.CyclicDefinitionError):
+            class TestObjectFactory(factory.Factory):
+                class Meta:
+                    model = TestObject
+
+                class Params:
+                    a = factory.Trait(b=True, one=True)
+                    b = factory.Trait(a=True, two=True)
+
+
 class SubFactoryTestCase(unittest.TestCase):
     def test_sub_factory(self):
         class TestModel2(FakeModel):