Add support for 'LazyContainerAttribute', when a SubFactory field needs access to attributes from its parent.

Signed-off-by: Raphaël Barrois <raphael.barrois@polytechnique.org>

1 files changed, 43 insertions, 6 deletions

diff --git a/factory/declarations.py b/factory/declarations.py
index 4a5bf97..c28e6af 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -28,7 +28,7 @@ class OrderedDeclaration(object):
     in the same factory.
     """
 
-    def evaluate(self, sequence, obj):
+    def evaluate(self, sequence, obj, containers=()):
         """Evaluate this declaration.
 
         Args:
@@ -36,6 +36,8 @@ class OrderedDeclaration(object):
                 the current instance
             obj (containers.LazyStub): The object holding currently computed
                 attributes
+            containers (list of containers.LazyStub): The chain of SubFactory
+                which led to building this object.
         """
         raise NotImplementedError('This is an abstract method')
 
@@ -52,7 +54,7 @@ class LazyAttribute(OrderedDeclaration):
         super(LazyAttribute, self).__init__(*args, **kwargs)
         self.function = function
 
-    def evaluate(self, sequence, obj):
+    def evaluate(self, sequence, obj, containers=()):
         return self.function(obj)
 
 
@@ -67,7 +69,7 @@ class SelfAttribute(OrderedDeclaration):
         super(SelfAttribute, self).__init__(*args, **kwargs)
         self.attribute_name = attribute_name
 
-    def evaluate(self, sequence, obj):
+    def evaluate(self, sequence, obj, containers=()):
         # TODO(rbarrois): allow the use of ATTR_SPLITTER to fetch fields of
         # subfactories.
         return getattr(obj, self.attribute_name)
@@ -89,7 +91,7 @@ class Sequence(OrderedDeclaration):
         self.function = function
         self.type = type
 
-    def evaluate(self, sequence, obj):
+    def evaluate(self, sequence, obj, containers=()):
         return self.function(self.type(sequence))
 
 
@@ -102,10 +104,41 @@ class LazyAttributeSequence(Sequence):
         type (function): A function converting an integer into the expected kind
             of counter for the 'function' attribute.
     """
-    def evaluate(self, sequence, obj):
+    def evaluate(self, sequence, obj, containers=()):
         return self.function(obj, self.type(sequence))
 
 
+class LazyContainerAttribute(OrderedDeclaration):
+    """Variant of LazyAttribute, also receives the containers of the object.
+
+    Attributes:
+        function (function): A function, expecting the current LazyStub and the
+            (optional) object having a subfactory containing this attribute.
+        strict (bool): Whether evaluating should fail when the containers are
+            not passed in (i.e used outside a SubFactory).
+    """
+    def __init__(self, function, strict=True, *args, **kwargs):
+        super(LazyContainerAttribute, self).__init__(*args, **kwargs)
+        self.function = function
+        self.strict = strict
+
+    def evaluate(self, sequence, obj, containers=()):
+        """Evaluate the current LazyContainerAttribute.
+
+        Args:
+            obj (LazyStub): a lazy stub of the object being constructed, if
+                needed.
+            containers (LazyStub): a lazy stub of a factory being evaluated, with
+                a SubFactory building 'obj'.
+        """
+        if self.strict and not containers:
+            raise TypeError(
+                "A LazyContainerAttribute in 'strict' mode can only be used "
+                "within a SubFactory.")
+
+        return self.function(obj, containers)
+
+
 class SubFactory(OrderedDeclaration):
     """Base class for attributes based upon a sub-factory.
 
@@ -120,7 +153,7 @@ class SubFactory(OrderedDeclaration):
         self.defaults = kwargs
         self.factory = factory
 
-    def evaluate(self, create, extra):
+    def evaluate(self, create, extra, containers):
         """Evaluate the current definition and fill its attributes.
 
         Uses attributes definition in the following order:
@@ -132,6 +165,7 @@ class SubFactory(OrderedDeclaration):
         defaults = dict(self.defaults)
         if extra:
             defaults.update(extra)
+        defaults['__containers'] = containers
 
         attrs = self.factory.attributes(create, defaults)
 
@@ -151,3 +185,6 @@ def sequence(func):
 
 def lazy_attribute_sequence(func):
     return LazyAttributeSequence(func)
+
+def lazy_container_attribute(func):
+    return LazyContainerAttribute(func, strict=False)