Add a new type of attribute, SelfAttribute.

Should be used as a shortcut for 'lambda o: o.some_attr' ; most useful in SubFactories.

Signed-off-by: Raphaël Barrois <raphael.barrois@polyconseil.fr>

2 files changed, 18 insertions, 0 deletions

diff --git a/factory/declarations.py b/factory/declarations.py
index 238e592..d827f30 100644
--- a/factory/declarations.py
+++ b/factory/declarations.py
@@ -62,6 +62,16 @@ class LazyAttribute(OrderedDeclaration):
     def evaluate(self, factory, attributes):
         return self.function(attributes)
 
+
+class SelfAttribute(OrderedDeclaration):
+    def __init__(self, attribute_name):
+        super(SelfAttribute, self).__init__()
+        self.attribute_name = attribute_name
+
+    def evaluate(self, factory, attributes):
+        return getattr(attributes, self.attribute_name)
+
+
 class Sequence(OrderedDeclaration):
     def __init__(self, function, type=str):
         super(Sequence, self).__init__()
diff --git a/factory/test_base.py b/factory/test_base.py
index 689b56d..b324680 100644
--- a/factory/test_base.py
+++ b/factory/test_base.py
@@ -115,6 +115,14 @@ class FactoryTestCase(unittest.TestCase):
         test_object = TestObjectFactory.build()
         self.assertEqual(test_object.one, 'one')
 
+    def testSelfAttribute(self):
+        class TestObjectFactory(Factory):
+            one = 'xx'
+            two = declarations.SelfAttribute('one')
+
+        test_object = TestObjectFactory.build(one=1)
+        self.assertEqual(1, test_object.two)
+
     def testSequenceDecorator(self):
         class TestObjectFactory(Factory):
             @declarations.sequence