Add PostGenerationDeclarationDict (cf. DeclarationDict).

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

1 files changed, 9 insertions, 1 deletions

diff --git a/factory/containers.py b/factory/containers.py
index b8557d6..9f480cc 100644
--- a/factory/containers.py
+++ b/factory/containers.py
@@ -149,13 +149,21 @@ class DeclarationDict(dict):
         Args:
             extra (dict): additional attributes to include in the copy.
         """
-        new = DeclarationDict()
+        new = self.__class__()
         new.update(self)
         if extra:
             new.update(extra)
         return new
 
 
+class PostGenerationDeclarationDict(DeclarationDict):
+    """Alternate DeclarationDict for PostGenerationDeclaration."""
+
+    def is_declaration(self, name, value):
+        """Captures instances of PostGenerationDeclaration."""
+        return isinstance(value, declarations.PostGenerationDeclaration)
+
+
 class LazyValue(object):
     """Some kind of "lazy evaluating" object."""