Enable models to be registered for multiple types

This addresses the underlying concern on #22, though in a different way
than is suggested there. The common case is registering for a single
type, and that is what MyModel.typname is intended for. If you want to
get fancy then use register_composite with a typname kwarg.

Author: chadwhitacre

postgres/__init__.py

@@ -599,11 +599,17 @@ def get_connection(self):
         return ConnectionContextManager(self.pool)
 
 
-    def register_model(self, ModelSubclass):
+    def register_model(self, ModelSubclass, typname=None):
         """Register an ORM model.
 
         :param ModelSubclass: the :py:class:`~postgres.orm.Model` subclass to
             register with this :py:class:`~postgres.Postgres` instance
+
+        :param typname: a string indicating the Postgres type to register this
+            model for (``typname``, without an "e," is the name of the relevant
+            column in the underlying ``pg_type`` table). If :py:class:`None`,
+            we'll look for :py:attr:`ModelSubclass.typname`.
+
         :raises: :py:exc:`~postgres.NotAModel`,
             :py:exc:`~postgres.NoTypeSpecified`,
             :py:exc:`~postgres.NoSuchType`,
@@ -618,28 +624,30 @@ def register_model(self, ModelSubclass):
         if not issubclass(ModelSubclass, Model):
             raise NotAModel(ModelSubclass)
 
-        if getattr(ModelSubclass, 'typname', None) is None:
-            raise NoTypeSpecified(ModelSubclass)
+        if typname is None:
+            typname = getattr(ModelSubclass, 'typname', None)
+            if typname is None:
+                raise NoTypeSpecified(ModelSubclass)
 
         n = self.one( "SELECT count(*) FROM pg_type WHERE typname=%s"
-                    , (ModelSubclass.typname,)
+                    , (typname,)
                      )
         if n < 1:
             # Could be more than one since we don't constrain by typnamespace.
             # XXX What happens then?
-            raise NoSuchType(ModelSubclass.typname)
+            raise NoSuchType(typname)
 
-        if ModelSubclass.typname in self.model_registry:
-            existing_model = self.model_registry[ModelSubclass.typname]
-            raise AlreadyRegistered(existing_model, ModelSubclass.typname)
+        if typname in self.model_registry:
+            existing_model = self.model_registry[typname]
+            raise AlreadyRegistered(existing_model, typname)
 
-        self.model_registry[ModelSubclass.typname] = ModelSubclass
+        self.model_registry[typname] = ModelSubclass
         ModelSubclass.db = self
 
-        # register a composite (but don't use RealDictCursor, not sure why)
+        # register a composite
         with self.get_connection() as conn:
             cursor = conn.cursor()
-            name = ModelSubclass.typname
+            name = typname
             if sys.version_info[0] < 3:
                 name = name.encode('UTF-8')
             register_composite( name
@@ -656,30 +664,40 @@ def unregister_model(self, ModelSubclass):
             unregister
         :raises: :py:exc:`~postgres.NotRegistered`
 
+        If :py:class:`ModelSubclass` is registered for multiple types, it is
+        unregistered for all of them.
+
         """
-        key = self.check_registration(ModelSubclass)
-        del self.model_registry[key]
+        keys = self.check_registration(ModelSubclass)
+        if not isinstance(keys, list):
+            # Wrap single string in a list. Flip-side of XXX just below.
+            keys = [keys]
+        for key in keys:
+            del self.model_registry[key]
 
 
     def check_registration(self, ModelSubclass):
         """Check whether an ORM model is registered.
 
         :param ModelSubclass: the :py:class:`~postgres.orm.Model` subclass to
             check for
+
         :returns: the :py:attr:`typname` (a string) for which this model is
-            registered
+            registered, or a list of strings if it's registered for multiple
+            types
+
         :rettype: string
         :raises: :py:exc:`~postgres.NotRegistered`
 
         """
-        key = None
-        for k, v in self.model_registry.items():
-            if v is ModelSubclass:
-                key = k
-                break
-        if key is None:
+        keys = [k for k,v in self.model_registry.items() if v is ModelSubclass]
+        if not keys:
             raise NotRegistered(ModelSubclass)
-        return key
+        if len(keys) == 1:
+            # Dereference a single-item list, for backwards-compatibility.
+            # XXX If/when we go to 3.0, lose this cruft (always return list).
+            keys = keys[0]
+        return keys
 
 
 # Class Factories

tests.py

@@ -6,7 +6,7 @@
 
 from postgres import Postgres, NotRegistered
 from postgres.cursors import TooFew, TooMany, SimpleDictCursor
-from postgres.orm import ReadOnly
+from postgres.orm import ReadOnly, Model
 from psycopg2 import InterfaceError, ProgrammingError
 from pytest import raises
 
@@ -223,14 +223,11 @@ def test_get_connection_gets_a_connection(self):
 
 class TestORM(WithData):
 
-    from postgres.orm import Model
-
     class MyModel(Model):
 
         typname = "foo"
 
         def __init__(self, record):
-            from postgres.orm import Model
             Model.__init__(self, record)
             self.bar_from_init = record['bar']
 
@@ -247,6 +244,10 @@ def setUp(self):
     def tearDown(self):
         self.db.model_registry = {}
 
+    def installFlah(self):
+        self.db.run("CREATE TABLE flah (bar text)")
+        self.db.register_model(self.MyModel, 'flah')
+
     def test_orm_basically_works(self):
         one = self.db.one("SELECT foo.*::foo FROM foo WHERE bar='baz'")
         assert one.__class__ == self.MyModel
@@ -271,6 +272,37 @@ def test_check_register_raises_if_passed_a_model_instance(self):
         obj = self.MyModel({'bar': 'baz'})
         raises(NotRegistered, self.db.check_registration, obj)
 
+    def test_a_model_can_be_used_for_a_second_type(self):
+        self.installFlah()
+        self.db.run("INSERT INTO flah VALUES ('double')")
+        self.db.run("INSERT INTO flah VALUES ('trouble')")
+        flah = self.db.one("SELECT flah.*::flah FROM flah WHERE bar='double'")
+        assert flah.bar == "double"
+
+    def test_check_register_returns_string_for_single(self):
+        assert self.db.check_registration(self.MyModel) == 'foo'
+
+    def test_check_register_returns_list_for_multiple(self):
+        self.installFlah()
+        actual = list(sorted(self.db.check_registration(self.MyModel)))
+        assert actual == ['flah', 'foo']
+
+    def test_unregister_unregisters_one(self):
+        self.db.unregister_model(self.MyModel)
+        assert self.db.model_registry == {}
+
+    def test_unregister_leaves_other(self):
+        self.db.run("CREATE TABLE flum (bar text)")
+        class OtherModel(Model): pass
+        self.db.register_model(OtherModel, 'flum')
+        self.db.unregister_model(self.MyModel)
+        assert self.db.model_registry == {'flum': OtherModel}
+
+    def test_unregister_unregisters_multiple(self):
+        self.installFlah()
+        self.db.unregister_model(self.MyModel)
+        assert self.db.model_registry == {}
+
 
 # cursor_factory
 # ==============