diff options
Diffstat (limited to 'lib/extras.py')
| -rw-r--r-- | lib/extras.py | 100 |
1 files changed, 63 insertions, 37 deletions
diff --git a/lib/extras.py b/lib/extras.py index 7a3a925..fe74d38 100644 --- a/lib/extras.py +++ b/lib/extras.py @@ -40,10 +40,23 @@ from psycopg2 import extensions as _ext from psycopg2.extensions import cursor as _cursor from psycopg2.extensions import connection as _connection from psycopg2.extensions import adapt as _A, quote_ident -from psycopg2._psycopg import REPLICATION_PHYSICAL, REPLICATION_LOGICAL -from psycopg2._psycopg import ReplicationConnection as _replicationConnection -from psycopg2._psycopg import ReplicationCursor as _replicationCursor -from psycopg2._psycopg import ReplicationMessage + +from psycopg2._psycopg import ( # noqa + REPLICATION_PHYSICAL, REPLICATION_LOGICAL, + ReplicationConnection as _replicationConnection, + ReplicationCursor as _replicationCursor, + ReplicationMessage) + + +# expose the json adaptation stuff into the module +from psycopg2._json import ( # noqa + json, Json, register_json, register_default_json, register_default_jsonb) + + +# Expose range-related objects +from psycopg2._range import ( # noqa + Range, NumericRange, DateRange, DateTimeRange, DateTimeTZRange, + register_range, RangeAdapter, RangeCaster) class DictCursorBase(_cursor): @@ -109,6 +122,7 @@ class DictConnection(_connection): kwargs.setdefault('cursor_factory', DictCursor) return super(DictConnection, self).cursor(*args, **kwargs) + class DictCursor(DictCursorBase): """A cursor that keeps a list of column name -> index mappings.""" @@ -133,6 +147,7 @@ class DictCursor(DictCursorBase): self.index[self.description[i][0]] = i self._query_executed = 0 + class DictRow(list): """A row object that allow by-column-name access to data.""" @@ -195,10 +210,10 @@ class DictRow(list): # drop the crusty Py2 methods if _sys.version_info[0] > 2: - items = iteritems; del iteritems - keys = iterkeys; del iterkeys - values = itervalues; del itervalues - del has_key + items = iteritems # noqa + keys = iterkeys # noqa + values = itervalues # noqa + del iteritems, iterkeys, itervalues, has_key class RealDictConnection(_connection): @@ -207,6 +222,7 @@ class RealDictConnection(_connection): kwargs.setdefault('cursor_factory', RealDictCursor) return super(RealDictConnection, self).cursor(*args, **kwargs) + class RealDictCursor(DictCursorBase): """A cursor that uses a real dict as the base type for rows. @@ -236,6 +252,7 @@ class RealDictCursor(DictCursorBase): self.column_mapping.append(self.description[i][0]) self._query_executed = 0 + class RealDictRow(dict): """A `!dict` subclass representing a data record.""" @@ -268,6 +285,7 @@ class NamedTupleConnection(_connection): kwargs.setdefault('cursor_factory', NamedTupleCursor) return super(NamedTupleConnection, self).cursor(*args, **kwargs) + class NamedTupleCursor(_cursor): """A cursor that generates results as `~collections.namedtuple`. @@ -372,11 +390,13 @@ class LoggingConnection(_connection): def _logtofile(self, msg, curs): msg = self.filter(msg, curs) - if msg: self._logobj.write(msg + _os.linesep) + if msg: + self._logobj.write(msg + _os.linesep) def _logtologger(self, msg, curs): msg = self.filter(msg, curs) - if msg: self._logobj.debug(msg) + if msg: + self._logobj.debug(msg) def _check(self): if not hasattr(self, '_logobj'): @@ -388,6 +408,7 @@ class LoggingConnection(_connection): kwargs.setdefault('cursor_factory', LoggingCursor) return super(LoggingConnection, self).cursor(*args, **kwargs) + class LoggingCursor(_cursor): """A cursor that logs queries using its connection logging facilities.""" @@ -428,6 +449,7 @@ class MinTimeLoggingConnection(LoggingConnection): kwargs.setdefault('cursor_factory', MinTimeLoggingCursor) return LoggingConnection.cursor(self, *args, **kwargs) + class MinTimeLoggingCursor(LoggingCursor): """The cursor sub-class companion to `MinTimeLoggingConnection`.""" @@ -479,18 +501,23 @@ class ReplicationCursor(_replicationCursor): if slot_type == REPLICATION_LOGICAL: if output_plugin is None: - raise psycopg2.ProgrammingError("output plugin name is required to create logical replication slot") + raise psycopg2.ProgrammingError( + "output plugin name is required to create " + "logical replication slot") command += "LOGICAL %s" % quote_ident(output_plugin, self) elif slot_type == REPLICATION_PHYSICAL: if output_plugin is not None: - raise psycopg2.ProgrammingError("cannot specify output plugin name when creating physical replication slot") + raise psycopg2.ProgrammingError( + "cannot specify output plugin name when creating " + "physical replication slot") command += "PHYSICAL" else: - raise psycopg2.ProgrammingError("unrecognized replication type: %s" % repr(slot_type)) + raise psycopg2.ProgrammingError( + "unrecognized replication type: %s" % repr(slot_type)) self.execute(command) @@ -513,7 +540,8 @@ class ReplicationCursor(_replicationCursor): if slot_name: command += "SLOT %s " % quote_ident(slot_name, self) else: - raise psycopg2.ProgrammingError("slot name is required for logical replication") + raise psycopg2.ProgrammingError( + "slot name is required for logical replication") command += "LOGICAL " @@ -523,28 +551,32 @@ class ReplicationCursor(_replicationCursor): # don't add "PHYSICAL", before 9.4 it was just START_REPLICATION XXX/XXX else: - raise psycopg2.ProgrammingError("unrecognized replication type: %s" % repr(slot_type)) + raise psycopg2.ProgrammingError( + "unrecognized replication type: %s" % repr(slot_type)) if type(start_lsn) is str: lsn = start_lsn.split('/') lsn = "%X/%08X" % (int(lsn[0], 16), int(lsn[1], 16)) else: - lsn = "%X/%08X" % ((start_lsn >> 32) & 0xFFFFFFFF, start_lsn & 0xFFFFFFFF) + lsn = "%X/%08X" % ((start_lsn >> 32) & 0xFFFFFFFF, + start_lsn & 0xFFFFFFFF) command += lsn if timeline != 0: if slot_type == REPLICATION_LOGICAL: - raise psycopg2.ProgrammingError("cannot specify timeline for logical replication") + raise psycopg2.ProgrammingError( + "cannot specify timeline for logical replication") command += " TIMELINE %d" % timeline if options: if slot_type == REPLICATION_PHYSICAL: - raise psycopg2.ProgrammingError("cannot specify output plugin options for physical replication") + raise psycopg2.ProgrammingError( + "cannot specify output plugin options for physical replication") command += " (" - for k,v in options.iteritems(): + for k, v in options.iteritems(): if not command.endswith('('): command += ", " command += "%s %s" % (quote_ident(k, self), _A(str(v))) @@ -579,6 +611,7 @@ class UUID_adapter(object): def __str__(self): return "'%s'::uuid" % self._uuid + def register_uuid(oids=None, conn_or_curs=None): """Create the UUID type and an uuid.UUID adapter. @@ -643,6 +676,7 @@ class Inet(object): def __str__(self): return str(self.addr) + def register_inet(oid=None, conn_or_curs=None): """Create the INET type and an Inet adapter. @@ -862,8 +896,9 @@ WHERE typname = 'hstore'; return tuple(rv0), tuple(rv1) + def register_hstore(conn_or_curs, globally=False, unicode=False, - oid=None, array_oid=None): + oid=None, array_oid=None): """Register adapter and typecaster for `!dict`\-\ |hstore| conversions. :param conn_or_curs: a connection or cursor: the typecaster will be @@ -942,8 +977,8 @@ class CompositeCaster(object): self.oid = oid self.array_oid = array_oid - self.attnames = [ a[0] for a in attrs ] - self.atttypes = [ a[1] for a in attrs ] + self.attnames = [a[0] for a in attrs] + self.atttypes = [a[1] for a in attrs] self._create_type(name, self.attnames) self.typecaster = _ext.new_type((oid,), name, self.parse) if array_oid: @@ -962,8 +997,8 @@ class CompositeCaster(object): "expecting %d components for the type %s, %d found instead" % (len(self.atttypes), self.name, len(tokens))) - values = [ curs.cast(oid, token) - for oid, token in zip(self.atttypes, tokens) ] + values = [curs.cast(oid, token) + for oid, token in zip(self.atttypes, tokens)] return self.make(values) @@ -1057,11 +1092,12 @@ ORDER BY attnum; type_oid = recs[0][0] array_oid = recs[0][1] - type_attrs = [ (r[2], r[3]) for r in recs ] + type_attrs = [(r[2], r[3]) for r in recs] return self(tname, type_oid, type_attrs, array_oid=array_oid, schema=schema) + def register_composite(name, conn_or_curs, globally=False, factory=None): """Register a typecaster to convert a composite type into a tuple. @@ -1084,17 +1120,7 @@ def register_composite(name, conn_or_curs, globally=False, factory=None): _ext.register_type(caster.typecaster, not globally and conn_or_curs or None) if caster.array_typecaster is not None: - _ext.register_type(caster.array_typecaster, not globally and conn_or_curs or None) + _ext.register_type( + caster.array_typecaster, not globally and conn_or_curs or None) return caster - - -# expose the json adaptation stuff into the module -from psycopg2._json import json, Json, register_json -from psycopg2._json import register_default_json, register_default_jsonb - - -# Expose range-related objects -from psycopg2._range import Range, NumericRange -from psycopg2._range import DateRange, DateTimeRange, DateTimeTZRange -from psycopg2._range import register_range, RangeAdapter, RangeCaster |