From 03c40fd80707ad4837523a07cdf3f82564ab0259 Mon Sep 17 00:00:00 2001 From: Raphaƫl Barrois Date: Sat, 2 Apr 2016 16:14:06 +0200 Subject: 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. --- factory/__init__.py | 1 + factory/declarations.py | 16 ++++++++++++++++ 2 files changed, 17 insertions(+) (limited to 'factory') 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 # =============== -- cgit v1.2.3