Relax check_registration to allow subsubclasses

This is necessary to allow __new__ in a Model subclass to return
subclasses of the subclass.

Author: chadwhitacre

postgres/__init__.py

@@ -168,6 +168,7 @@
 from collections import namedtuple
 
 import psycopg2
+from inspect import isclass
 from postgres.context_managers import ConnectionContextManager
 from postgres.context_managers import CursorContextManager
 from postgres.cursors import SimpleTupleCursor, SimpleNamedTupleCursor
@@ -214,8 +215,8 @@ def __str__(self):
 class NotAModel(Exception):
     def __str__(self):
         return "Only subclasses of postgres.orm.Model can be registered as " \
-               "orm models. {} (registered for {}) doesn't fit the bill." \
-               .format(self.args[0].__name__, self.args[1])
+               "orm models. {} doesn't fit the bill." \
+               .format(self.args[0])
 
 class NoTypeSpecified(Exception):
     def __str__(self):
@@ -621,8 +622,7 @@ def register_model(self, ModelSubclass, typname=None):
             subclassing :py:class:`~postgres.orm.Model`.
 
         """
-        if not issubclass(ModelSubclass, Model):
-            raise NotAModel(ModelSubclass)
+        self._validate_model_subclass(ModelSubclass)
 
         if typname is None:
             typname = getattr(ModelSubclass, 'typname', None)
@@ -676,11 +676,14 @@ def unregister_model(self, ModelSubclass):
             del self.model_registry[key]
 
 
-    def check_registration(self, ModelSubclass):
+    def check_registration(self, ModelSubclass, include_subsubclasses=False):
         """Check whether an ORM model is registered.
 
         :param ModelSubclass: the :py:class:`~postgres.orm.Model` subclass to
             check for
+        :param bool include_subsubclasses: whether to also check for subclasses
+            of :py:class:`ModelSubclass` or just :py:class:`ModelSubclass`
+            itself
 
         :returns: the :py:attr:`typname` (a string) for which this model is
             registered, or a list of strings if it's registered for multiple
@@ -690,7 +693,13 @@ def check_registration(self, ModelSubclass):
         :raises: :py:exc:`~postgres.NotRegistered`
 
         """
-        keys = [k for k,v in self.model_registry.items() if v is ModelSubclass]
+        self._validate_model_subclass(ModelSubclass)
+
+        if include_subsubclasses:
+            filt = lambda v: v is ModelSubclass or issubclass(ModelSubclass, v)
+        else:
+            filt = lambda v: v is ModelSubclass
+        keys = [k for k,v in self.model_registry.items() if filt(v)]
         if not keys:
             raise NotRegistered(ModelSubclass)
         if len(keys) == 1:
@@ -700,6 +709,11 @@ def check_registration(self, ModelSubclass):
         return keys
 
 
+    def _validate_model_subclass(self, ModelSubclass, ):
+        if not isclass(ModelSubclass) or not issubclass(ModelSubclass, Model):
+            raise NotAModel(ModelSubclass)
+
+
 # Class Factories
 # ===============
 

postgres/orm.py

@@ -223,7 +223,7 @@ class Model(object):
     def __init__(self, record):
         if self.db is None:
             raise NotBound(self)
-        self.db.check_registration(self.__class__)
+        self.db.check_registration(self.__class__, include_subsubclasses=True)
         self.__read_only_attributes = record.keys()
         self.set_attributes(**record)
 

tests.py

@@ -4,7 +4,7 @@
 from collections import namedtuple
 from unittest import TestCase
 
-from postgres import Postgres, NotRegistered
+from postgres import Postgres, NotAModel, NotRegistered
 from postgres.cursors import TooFew, TooMany, SimpleDictCursor
 from postgres.orm import ReadOnly, Model
 from psycopg2 import InterfaceError, ProgrammingError
@@ -270,7 +270,27 @@ def assign():
 
     def test_check_register_raises_if_passed_a_model_instance(self):
         obj = self.MyModel({'bar': 'baz'})
-        raises(NotRegistered, self.db.check_registration, obj)
+        raises(NotAModel, self.db.check_registration, obj)
+
+    def test_check_register_doesnt_include_subsubclasses(self):
+        class Other(self.MyModel): pass
+        raises(NotRegistered, self.db.check_registration, Other)
+
+    def test_dot_dot_dot_unless_you_ask_it_to(self):
+        class Other(self.MyModel): pass
+        assert self.db.check_registration(Other, True) == 'foo'
+
+    def test_check_register_handles_complex_cases(self):
+        self.installFlah()
+
+        class Second(Model): pass
+        self.db.run("CREATE TABLE blum (bar text)")
+        self.db.register_model(Second, 'blum')
+        assert self.db.check_registration(Second) == 'blum'
+
+        class Third(self.MyModel, Second): pass
+        actual = list(sorted(self.db.check_registration(Third, True)))
+        assert actual == ['blum', 'flah', 'foo']
 
     def test_a_model_can_be_used_for_a_second_type(self):
         self.installFlah()