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.

2 files changed, 17 insertions, 0 deletions

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
 # ===============