Clean up back_as

We were calling handle_back_as in two places where we really only needed
one.

Author: chadwhitacre

postgres/__init__.py

@@ -165,12 +165,13 @@
 
 except ImportError:     # Python 3
     import urllib.parse as urlparse
+from collections import namedtuple
 
 import psycopg2
 from postgres.context_managers import ConnectionContextManager
 from postgres.context_managers import CursorContextManager
-from postgres.context_managers import handle_back_as
-from postgres.cursors import SimpleNamedTupleCursor, SimpleCursorBase
+from postgres.cursors import SimpleTupleCursor, SimpleNamedTupleCursor
+from postgres.cursors import SimpleDictCursor, SimpleCursorBase
 from postgres.orm import Model
 from psycopg2.extras import register_composite, CompositeCaster
 from psycopg2.pool import ThreadedConnectionPool as ConnectionPool
@@ -236,6 +237,12 @@ class NotRegistered(Exception):
     def __str__(self):
         return "The model {} is not registered.".format(self.args[0].__name__)
 
+class BadBackAs(Exception):
+    def __str__(self):
+        return "Bad back_as: {}. Available back_as values are: tuple, " \
+               "namedtuple, dict, or None (to use the default)." \
+               .format(self.args[0])
+
 
 # The Main Event
 # ==============
@@ -522,10 +529,15 @@ def all(self, sql, parameters=None, back_as=None, *a, **kw):
             return cursor.all(sql, parameters)
 
 
-    def get_cursor(self, back_as=None, *a, **kw):
+    def get_cursor(self, *a, **kw):
         """Return a :py:class:`~postgres.CursorContextManager` that uses
         our connection pool.
 
+        :param a: passed through to the :py:meth:`cursor` method of instances
+            of the class returned by :py:func:`~postgres.make_Connection`
+        :param kw: passed through to the :py:meth:`cursor` method of instances
+            of the class returned by :py:func:`~postgres.make_Connection`
+
         >>> with db.get_cursor() as cursor:
         ...     cursor.all("SELECT * FROM foo")
         ...
@@ -554,7 +566,7 @@ def get_cursor(self, back_as=None, *a, **kw):
         transaction.
 
         """
-        return CursorContextManager(self.pool, back_as=back_as, *a, **kw)
+        return CursorContextManager(self.pool, *a, **kw)
 
 
     def get_connection(self):
@@ -680,14 +692,26 @@ def make_Connection(postgres):
     :returns: a :py:class:`Connection` class
 
     The class defined and returned here will be linked to the instance of
-    :py:class:`~postgres.Postgres` that is passed in as :py:attr:`postgres` and
-    will use the :py:attr:`default_cursor_factory` attribute of that object.
-    The :py:class:`~postgres.Postgres` instance will use this class as the
-    :py:attr:`connection_factory` for its connection pool. The
-    :py:meth:`cursor` method of this class accepts a :py:attr:`back_as`
-    argument, which is processed according to
-    :py:func:`~postgres.handle_back_as`. We also set client encoding to
-    ``UTF-8``.
+    :py:class:`~postgres.Postgres` that is passed in as :py:attr:`postgres`,
+    which will use this class as the :py:attr:`connection_factory` for its
+    connection pool.
+
+    The :py:meth:`cursor` method of this class accepts a :py:attr:`back_as`
+    keyword argument. If a :py:attr:`cursor_factory` keyword argument is also
+    given, then any :py:attr:`back_as` is ignored and discarded.  Valid values
+    for :py:attr:`back_as` are :py:class:`tuple`, :py:class:`namedtuple`,
+    :py:class:`dict` (or the strings ``tuple``, ``namedtuple``, and ``dict``),
+    and :py:class:`None`. If the value of :py:attr:`back_as` is
+    :py:class:`None`, then we'll use the default :py:attr:`cursor_factory` with
+    which our parent :py:class:`~postgres.Postgres` instance was instantiated.
+    If :py:attr:`back_as` is not :py:class:`None`, then we'll specify a
+    :py:attr:`cursor_factory` that will result in records of the designated
+    type: :py:class:`postgres.cursor.SimpleTupleCursor` for :py:class:`tuple`,
+    :py:class:`postgres.cursor.SimpleNamedTupleCursor` for
+    :py:class:`namedtuple`, and :py:class:`postgres.cursor.SimpleDictCursor`
+    for :py:class:`dict`.
+
+    We also set client encoding to ``UTF-8``.
 
     """
     class Connection(psycopg2.extensions.connection):
@@ -700,11 +724,36 @@ def __init__(self, *a, **kw):
         def cursor(self, *a, **kw):
             if 'back_as' in kw:
                 back_as = kw.pop('back_as')
-                kw = handle_back_as(back_as, **kw)
+                kw = self.handle_back_as(back_as, **kw)
             if 'cursor_factory' not in kw:
                 kw['cursor_factory'] = self.postgres.default_cursor_factory
             return psycopg2.extensions.connection.cursor(self, *a, **kw)
 
+        def handle_back_as(self, back_as, **kw):
+
+            if 'cursor_factory' not in kw:
+
+                # Compute cursor_factory from back_as.
+                # ====================================
+
+                registry = { tuple: SimpleTupleCursor
+                           , 'tuple': SimpleTupleCursor
+                           , namedtuple: SimpleNamedTupleCursor
+                           , 'namedtuple': SimpleNamedTupleCursor
+                           , dict: SimpleDictCursor
+                           , 'dict': SimpleDictCursor
+                           , None: None
+                             }
+
+                if back_as not in registry:
+                    raise BadBackAs(back_as)
+
+                cursor_factory = registry[back_as]
+                if cursor_factory is not None:
+                    kw['cursor_factory'] = cursor_factory
+
+            return kw
+
     return Connection
 
 

postgres/context_managers.py

@@ -1,17 +1,5 @@
 from __future__ import print_function, unicode_literals
 
-from collections import namedtuple
-
-from postgres.cursors import SimpleTupleCursor, SimpleNamedTupleCursor
-from postgres.cursors import SimpleDictCursor
-
-
-class BadRecordType(Exception):
-    def __str__(self):
-        return "Bad back_as: {}. Available back_as values are: tuple, " \
-               "namedtuple, dict, or None (to use the default)." \
-               .format(self.args[0])
-
 
 class CursorContextManager(object):
     """Instantiated once per :py:func:`~postgres.Postgres.get_cursor`
@@ -21,10 +9,7 @@ class CursorContextManager(object):
 
     The return value of :py:func:`CursorContextManager.__enter__` is a
     :py:mod:`psycopg2` cursor. Any positional and keyword arguments to our
-    constructor are passed through to the cursor constructor. If you pass
-    :py:attr:`back_as` as a keyword argument then we'll infer a
-    :py:attr:`cursor_factory` from that, though any explicit
-    :py:attr:`cursor_factory` keyword argument will take precedence.
+    constructor are passed through to the cursor constructor.
 
     When the block starts, a connection is checked out of the connection pool
     and :py:attr:`autocommit` is set to :py:const:`False`. Then a cursor is
@@ -38,10 +23,10 @@ class CursorContextManager(object):
 
     """
 
-    def __init__(self, pool, back_as=None, *a, **kw):
+    def __init__(self, pool, *a, **kw):
         self.pool = pool
         self.a = a
-        self.kw = handle_back_as(back_as, **kw)
+        self.kw = kw
         self.conn = None
 
     def __enter__(self):
@@ -95,43 +80,3 @@ def __exit__(self, *exc_info):
         self.conn.rollback()
         self.conn.autocommit = False
         self.pool.putconn(self.conn)
-
-
-def handle_back_as(back_as, **kw):
-    """Add :py:attr:`cursor_factory` to :py:attr:`kw`, maybe.
-
-    Valid values for :py:attr:`back_as` are :py:class:`tuple`,
-    :py:class:`namedtuple`, :py:class:`dict` (or the strings ``tuple``,
-    ``namedtuple``, and ``dict``), and :py:class:`None`. If the value of
-    :py:attr:`back_as` is :py:class:`None`, then we won't insert any
-    :py:attr:`cursor_factory` keyword argument. Otherwise we'll specify a
-    :py:attr:`cursor_factory` that will result in records of the designated
-    type: :py:class:`postgres.cursor.SimpleTupleCursor` for :py:class:`tuple`,
-    :py:class:`postgres.cursor.SimpleNamedTupleCursor` for
-    :py:class:`namedtuple`, and :py:class:`postgres.cursor.SimpleDictCursor`
-    for :py:class:`dict`.
-
-    """
-
-    if 'cursor_factory' not in kw:
-
-        # Compute cursor_factory from back_as.
-        # ====================================
-
-        cursor_factory_registry = { tuple: SimpleTupleCursor
-                                  , 'tuple': SimpleTupleCursor
-                                  , namedtuple: SimpleNamedTupleCursor
-                                  , 'namedtuple': SimpleNamedTupleCursor
-                                  , dict: SimpleDictCursor
-                                  , 'dict': SimpleDictCursor
-                                  , None: None
-                                    }
-
-        if back_as not in cursor_factory_registry:
-            raise BadRecordType(back_as)
-
-        cursor_factory = cursor_factory_registry[back_as]
-        if cursor_factory is not None:
-            kw['cursor_factory'] = cursor_factory
-
-    return kw

postgres/cursors.py

@@ -9,7 +9,7 @@
 
 from inspect import isclass
 
-from psycopg2.extensions import cursor as RegularCursor
+from psycopg2.extensions import cursor as TupleCursor
 from psycopg2.extras import NamedTupleCursor, RealDictCursor
 
 
@@ -49,7 +49,7 @@ class SimpleCursorBase(object):
     :py:class:`~postgres.Postgres` must subclass this base.
 
     >>> from psycopg2.extras import LoggingCursor
-    >>> from postgres.cursor import SimpleCursorBase
+    >>> from postgres.cursors import SimpleCursorBase
     >>> class SimpleLoggingCursor(LoggingCursor, SimpleCursorBase):
     ...     pass
     ...
@@ -59,7 +59,7 @@ class SimpleCursorBase(object):
     ...               )
 
     If you try to use a cursor that doesn't subclass
-    :py:class:`~postgres.cursor.SimpleCursorBase` as the default
+    :py:class:`~postgres.cursors.SimpleCursorBase` as the default
     :py:attr:`cursor_factory` for a :py:class:`~postgres.Postgres` instance, we
     won't let you:
 
@@ -69,7 +69,7 @@ class SimpleCursorBase(object):
     ...
     Traceback (most recent call last):
         ...
-    postgres.NotASimpleCursor: We can only work with subclasses of postgres.cursor.SimpleCursorBase. LoggingCursor doesn't fit the bill.
+    postgres.NotASimpleCursor: We can only work with subclasses of postgres.cursors.SimpleCursorBase. LoggingCursor doesn't fit the bill.
 
     However, we do allow you to use whatever you want as the
     :py:attr:`cursor_factory` argument for individual calls:
@@ -152,7 +152,7 @@ def all(self, sql, parameters=None):
         return recs
 
 
-class SimpleTupleCursor(RegularCursor, SimpleCursorBase):
+class SimpleTupleCursor(TupleCursor, SimpleCursorBase):
     """A simple cursor that returns tuples.
     """