Cleaner code for ordered attributes.

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

2 files changed, 59 insertions, 13 deletions

diff --git a/factory/base.py b/factory/base.py
index 222c1bf..013b58f 100644
--- a/factory/base.py
+++ b/factory/base.py
@@ -21,7 +21,7 @@
 import re
 import sys
 
-from containers import ObjectParamsWrapper, StubObject
+from containers import ObjectParamsWrapper, OrderedDeclarationDict, StubObject
 from declarations import OrderedDeclaration
 
 # Strategies
@@ -78,20 +78,15 @@ class BaseFactoryMetaClass(type):
             # If this isn't a subclass of Factory, don't do anything special.
             return super(BaseFactoryMetaClass, cls).__new__(cls, class_name, bases, attrs)
 
-        ordered_declarations = getattr(base, CLASS_ATTRIBUTE_ORDERED_DECLARATIONS, [])
-        unordered_declarations = getattr(base, CLASS_ATTRIBUTE_UNORDERED_DECLARATIONS, [])
+        ordered_declarations = getattr(base, CLASS_ATTRIBUTE_ORDERED_DECLARATIONS,
+                                       OrderedDeclarationDict())
+        unordered_declarations = getattr(base, CLASS_ATTRIBUTE_UNORDERED_DECLARATIONS, {})
 
         for name in list(attrs):
             if isinstance(attrs[name], OrderedDeclaration):
-                ordered_declarations = [(_name, declaration) for (_name, declaration) in ordered_declarations if _name != name]
-                ordered_declarations.append((name, attrs[name]))
-                del attrs[name]
+                ordered_declarations[name] = attrs.pop(name)
             elif not name.startswith('_'):
-                unordered_declarations = [(_name, value) for (_name, value) in unordered_declarations if _name != name]
-                unordered_declarations.append((name, attrs[name]))
-                del attrs[name]
-
-        ordered_declarations.sort(key=lambda d: d[1].order)
+                unordered_declarations[name] = attrs.pop(name)
 
         attrs[CLASS_ATTRIBUTE_ORDERED_DECLARATIONS] = ordered_declarations
         attrs[CLASS_ATTRIBUTE_UNORDERED_DECLARATIONS] = unordered_declarations
@@ -169,16 +164,25 @@ class BaseFactory(object):
 
     @classmethod
     def attributes(cls, **kwargs):
+        """Build a dict of attribute values, respecting declaration order.
+
+        The process is:
+        - Handle 'orderless' attributes, overriding defaults with provided
+            kwargs when applicable
+        - Handle ordered attributes, overriding them with provided kwargs when
+            applicable; the current list of computed attributes is available for
+            to the currently processed object.
+        """
         attributes = {}
         cls.sequence = cls._generate_next_sequence()
 
-        for name, value in getattr(cls, CLASS_ATTRIBUTE_UNORDERED_DECLARATIONS):
+        for name, value in getattr(cls, CLASS_ATTRIBUTE_UNORDERED_DECLARATIONS).iteritems():
             if name in kwargs:
                 attributes[name] = kwargs.pop(name)
             else:
                 attributes[name] = value
 
-        for name, ordered_declaration in getattr(cls, CLASS_ATTRIBUTE_ORDERED_DECLARATIONS):
+        for name, ordered_declaration in getattr(cls, CLASS_ATTRIBUTE_ORDERED_DECLARATIONS).iteritems():
             if name in kwargs:
                 attributes[name] = kwargs.pop(name)
             else:
diff --git a/factory/containers.py b/factory/containers.py
index 023055c..a117d5c 100644
--- a/factory/containers.py
+++ b/factory/containers.py
@@ -41,6 +41,48 @@ class ObjectParamsWrapper(object):
         except KeyError:
             raise AttributeError("The param '{0}' does not exist. Perhaps your declarations are out of order?".format(name))
 
+
+class OrderedDeclarationDict(object):
+    def __init__(self, **kwargs):
+        self._order = {}
+        self._values = {}
+        for k, v in kwargs.iteritems():
+            self[k] = v
+
+    def __contains__(self, key):
+        return key in self._values
+
+    def __getitem__(self, key):
+        return self._values[key]
+
+    def __setitem__(self, key, val):
+        if key in self:
+            del self[key]
+        self._values[key] = val
+        self._order[val.order] = key
+
+    def __delitem__(self, key):
+        self.pop(key)
+
+    def pop(self, key):
+        val = self._values.pop(key)
+        del self._order[val.order]
+        return val
+
+    def items(self):
+        return list(self.iteritems())
+
+    def iteritems(self):
+        order = sorted(self._order.keys())
+        for i in order:
+            key = self._order[i]
+            yield (key, self._values[key])
+
+    def __iter__(self):
+        order = sorted(self._order.keys())
+        for i in order:
+            yield self._order[i]
+
 class StubObject(object):
     '''A generic container.'''