diff options
Diffstat (limited to 'factory')
-rw-r--r-- | factory/__init__.py | 1 | ||||
-rw-r--r-- | factory/declarations.py | 16 |
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 # =============== |