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>

2 files changed, 65 insertions, 14 deletions

diff --git a/factory/containers.py b/factory/containers.py
index e280162..4127851 100644
--- a/factory/containers.py
+++ b/factory/containers.py
@@ -49,10 +49,11 @@ class LazyStub(object):
 
     __initialized = False
 
-    def __init__(self, attrs):
+    def __init__(self, attrs, containers=()):
         self.__attrs = attrs
         self.__values = {}
         self.__pending = []
+        self.__containers = containers
         self.__initialized = True
 
     def __fill__(self):
@@ -82,7 +83,7 @@ class LazyStub(object):
             val = self.__attrs[name]
             if isinstance(val, LazyValue):
                 self.__pending.append(name)
-                val = val.evaluate(self)
+                val = val.evaluate(self, self.__containers)
                 assert name == self.__pending.pop()
             self.__values[name] = val
             return val
@@ -135,7 +136,7 @@ class DeclarationDict(dict):
 class LazyValue(object):
     """Some kind of "lazy evaluating" object."""
 
-    def evaluate(self, obj):
+    def evaluate(self, obj, containers=()):
         """Compute the value, using the given object."""
         raise NotImplementedError("This is an abstract method.")
 
@@ -156,8 +157,12 @@ class SubFactoryWrapper(LazyValue):
         self.subfields = subfields
         self.create = create
 
-    def evaluate(self, obj):
-        return self.subfactory.evaluate(self.create, self.subfields)
+    def evaluate(self, obj, containers=()):
+        expanded_containers = (obj,)
+        if containers:
+            expanded_containers += tuple(containers)
+        return self.subfactory.evaluate(self.create, self.subfields,
+            expanded_containers)
 
 
 class OrderedDeclarationWrapper(LazyValue):
@@ -175,8 +180,15 @@ class OrderedDeclarationWrapper(LazyValue):
         self.declaration = declaration
         self.sequence = sequence
 
-    def evaluate(self, obj):
-        return self.declaration.evaluate(self.sequence, obj)
+    def evaluate(self, obj, containers=()):
+        """Lazily evaluate the attached OrderedDeclaration.
+
+        Args:
+            obj (LazyStub): the object being built
+            containers (object list): the chain of containers of the object
+                being built, its immediate holder being first.
+        """
+        return self.declaration.evaluate(self.sequence, obj, containers)
 
 
 class AttributeBuilder(object):
@@ -195,7 +207,9 @@ class AttributeBuilder(object):
 
         if not extra:
             extra = {}
+
         self.factory = factory
+        self._containers = extra.pop('__containers', None)
         self._attrs = factory.declarations(extra)
         self._subfields = self._extract_subfields()
 
@@ -229,7 +243,7 @@ class AttributeBuilder(object):
                 v = OrderedDeclarationWrapper(v, self.factory.sequence)
             wrapped_attrs[k] = v
 
-        return LazyStub(wrapped_attrs).__fill__()
+        return LazyStub(wrapped_attrs, containers=self._containers).__fill__()
 
 
 class StubObject(object):
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)