allow easy override of Relation for a class

tpett · Sep 8, 2011 · 877b530 · 877b530
1 parent d34d805
commit 877b530
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/pyperry/base.py b/pyperry/base.py
@@ -386,6 +386,7 @@ class Article(pyperry.base.Base):
     """
 
     __metaclass__ = BaseMeta
+    _relation_class = Relation
 
     def __init__(self, fields=None, new_record=True, **kwargs):
         """
@@ -806,7 +807,7 @@ def relation(cls):
 
         """
         if not hasattr(cls, '_relation') or cls._relation.klass != cls:
-            cls._relation = Relation(cls)
+            cls._relation = cls._relation_class(cls)
         return cls._relation
 
     @classmethod

diff --git a/tests/base_test.py b/tests/base_test.py
@@ -879,6 +879,18 @@ def test_fresh(self):
         test.reload()
         self.assertEqual(test.last_relation.query()['fresh'], True)
 
+class RelationOverrideTestCase(BaseTestCase):
+
+    def test_uses_config_option(self):
+        class MyRelation(pyperry.relation.Relation): pass
+        class Test(pyperry.Base):
+            _relation_class = MyRelation
+
+        self.assertEqual(Test.relation().__class__, MyRelation)
+        self.assertEqual(
+                pyperry.base.Base.relation().__class__,
+                pyperry.relation.Relation)
+
 
 ##
 # Freezing objects