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>

4 files changed, 145 insertions, 15 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)
diff --git a/tests/test_base.py b/tests/test_base.py
index aaaa65f..9e61378 100644
--- a/tests/test_base.py
+++ b/tests/test_base.py
@@ -467,6 +467,60 @@ class FactoryDefaultStrategyTestCase(unittest.TestCase):
         self.assertEqual(wrapping.wrapped.two, 4)
         self.assertEqual(wrapping.friend, 5)
 
+    def testDiamondSubFactory(self):
+        """Tests the case where an object has two fields with a common field."""
+        class InnerMost(object):
+            def __init__(self, a, b):
+                self.a = a
+                self.b = b
+
+        class SideA(object):
+            def __init__(self, inner_from_a):
+                self.inner_from_a = inner_from_a
+
+        class SideB(object):
+            def __init__(self, inner_from_b):
+                self.inner_from_b = inner_from_b
+
+        class OuterMost(object):
+            def __init__(self, foo, side_a, side_b):
+                self.foo = foo
+                self.side_a = side_a
+                self.side_b = side_b
+
+        class InnerMostFactory(base.Factory):
+            FACTORY_FOR = InnerMost
+            a = 15
+            b = 20
+
+        class SideAFactory(base.Factory):
+            FACTORY_FOR = SideA
+            inner_from_a = declarations.SubFactory(InnerMostFactory, a=20)
+
+        class SideBFactory(base.Factory):
+            FACTORY_FOR = SideB
+            inner_from_b = declarations.SubFactory(InnerMostFactory, b=15)
+
+        class OuterMostFactory(base.Factory):
+            FACTORY_FOR = OuterMost
+
+            foo = 30
+            side_a = declarations.SubFactory(SideAFactory,
+                inner_from_a__a=declarations.LazyContainerAttribute(lambda obj, containers: containers[1].foo * 2))
+            side_b = declarations.SubFactory(SideBFactory,
+                inner_from_b=declarations.LazyContainerAttribute(lambda obj, containers: containers[0].side_a.inner_from_a))
+
+        outer = OuterMostFactory.build()
+        self.assertEqual(outer.foo, 30)
+        self.assertEqual(outer.side_a.inner_from_a, outer.side_b.inner_from_b)
+        self.assertEqual(outer.side_a.inner_from_a.a, outer.foo * 2)
+        self.assertEqual(outer.side_a.inner_from_a.b, 20)
+
+        outer = OuterMostFactory.build(side_a__inner_from_a__b = 4)
+        self.assertEqual(outer.foo, 30)
+        self.assertEqual(outer.side_a.inner_from_a, outer.side_b.inner_from_b)
+        self.assertEqual(outer.side_a.inner_from_a.a, outer.foo * 2)
+        self.assertEqual(outer.side_a.inner_from_a.b, 4)
 
     def testStubStrategy(self):
         base.Factory.default_strategy = base.STUB_STRATEGY
diff --git a/tests/test_containers.py b/tests/test_containers.py
index 34b61d4..57c06cf 100644
--- a/tests/test_containers.py
+++ b/tests/test_containers.py
@@ -43,12 +43,37 @@ class LazyStubTestCase(unittest.TestCase):
 
         self.assertRaises(AttributeError, getattr, stub, 'three')
 
+    def test_accessing_container(self):
+        class LazyAttr(containers.LazyValue):
+            def __init__(self, obj_attr, container_attr):
+                self.obj_attr = obj_attr
+                self.container_attr = container_attr
+
+            def evaluate(self, obj, containers=()):
+                if containers:
+                    add = getattr(containers[0], self.container_attr)
+                else:
+                    add = 0
+                return getattr(obj, self.obj_attr) + add
+
+        class DummyContainer(object):
+            three = 3
+
+        stub = containers.LazyStub({'one': LazyAttr('two', 'three'), 'two': 2, 'three': 42},
+            containers=(DummyContainer(),))
+
+        self.assertEqual(5, stub.one)
+
+        stub = containers.LazyStub({'one': LazyAttr('two', 'three'), 'two': 2, 'three': 42},
+            containers=())
+        self.assertEqual(2, stub.one)
+
     def test_cyclic_definition(self):
         class LazyAttr(containers.LazyValue):
             def __init__(self, attrname):
                 self.attrname = attrname
 
-            def evaluate(self, obj):
+            def evaluate(self, obj, container=None):
                 return 1 + getattr(obj, self.attrname)
 
         stub = containers.LazyStub({'one': LazyAttr('two'), 'two': LazyAttr('one')})