Changeset 770 for trunk/pg.py


Ignore:
Timestamp:
Jan 20, 2016, 1:19:45 PM (3 years ago)
Author:
cito
Message:

Add methods for getting a table as a list or dict

Also added documentation and 100% test coverage.

The get_attnames() method now always returns a read-only ordered dictionary,
even under Python 2.6 or 3.0. So you can sure the columns will be returned
in the right order if you iterate over it, and that you don't accidentally
modify the dictionary (since it is cached).

File:
1 edited

Legend:

Unmodified
Added
Removed
  • trunk/pg.py

    r769 r770  
    3838from collections import namedtuple
    3939from functools import partial
     40from operator import itemgetter
     41
     42try:
     43    basestring
     44except NameError:  # Python >= 3.0
     45    basestring = (str, bytes)
     46
     47set_decimal(Decimal)
    4048
    4149try:
     
    4452    OrderedDict = dict
    4553
    46 try:
    47     basestring
    48 except NameError:  # Python >= 3.0
    49     basestring = (str, bytes)
    50 
    51 set_decimal(Decimal)
     54
     55    class AttrDict(dict):
     56        """Simple read-only ordered dictionary for storing attribute names."""
     57
     58        def __init__(self, *args, **kw):
     59            if len(args) > 1 or kw:
     60                raise TypeError
     61            items = args[0] if args else []
     62            if isinstance(items, dict):
     63                raise TypeError
     64            items = list(items)
     65            self._keys = [item[0] for item in items]
     66            dict.__init__(self, items)
     67            self._read_only = True
     68            error = self._read_only_error
     69            self.clear = self.update = error
     70            self.pop = self.setdefault = self.popitem = error
     71
     72        def __setitem__(self, key, value):
     73            if self._read_only:
     74                self._read_only_error()
     75            dict.__setitem__(self, key, value)
     76
     77        def __delitem__(self, key):
     78            if self._read_only:
     79                self._read_only_error()
     80            dict.__delitem__(self, key)
     81
     82        def __iter__(self):
     83            return iter(self._keys)
     84
     85        def keys(self):
     86            return list(self._keys)
     87
     88        def values(self):
     89            return [self[key] for key in self]
     90
     91        def items(self):
     92            return [(key, self[key]) for key in self]
     93
     94        def iterkeys(self):
     95            return self.__iter__()
     96
     97        def itervalues(self):
     98            return iter(self.values())
     99
     100        def iteritems(self):
     101            return iter(self.items())
     102
     103        @staticmethod
     104        def _read_only_error(*args, **kw):
     105            raise TypeError('This object is read-only')
     106
     107else:
     108
     109     class AttrDict(OrderedDict):
     110        """Simple read-only ordered dictionary for storing attribute names."""
     111
     112        def __init__(self, *args, **kw):
     113            self._read_only = False
     114            OrderedDict.__init__(self, *args, **kw)
     115            self._read_only = True
     116            error = self._read_only_error
     117            self.clear = self.update = error
     118            self.pop = self.setdefault = self.popitem = error
     119
     120        def __setitem__(self, key, value):
     121            if self._read_only:
     122                self._read_only_error()
     123            OrderedDict.__setitem__(self, key, value)
     124
     125        def __delitem__(self, key):
     126            if self._read_only:
     127                self._read_only_error()
     128            OrderedDict.__delitem__(self, key)
     129
     130        @staticmethod
     131        def _read_only_error(*args, **kw):
     132            raise TypeError('This object is read-only')
     133
    52134
    53135
     
    84166
    85167set_namedresult(_namedresult)
     168
     169
     170class _MemoryQuery:
     171    """Class that embodies a given query result."""
     172
     173    def __init__(self, result, fields):
     174        """Create query from given result rows and field names."""
     175        self.result = result
     176        self.fields = fields
     177
     178    def listfields(self):
     179        """Return the stored field names of this query."""
     180        return self.fields
     181
     182    def getresult(self):
     183        """Return the stored result of this query."""
     184        return self.result
    86185
    87186
     
    704803        """Given the name of a table, dig out the set of attribute names.
    705804
    706         Returns a dictionary of attribute names (the names are the keys,
    707         the values are the names of the attributes' types).
    708 
    709         If your Python version supports this, the dictionary will be an
    710         OrderedDictionary with the column names in the right order.
    711 
    712         If flush is set then the internal cache for attribute names will
     805        Returns a read-only dictionary of attribute names (the names are
     806        the keys, the values are the names of the attributes' types)
     807        with the column names in the proper order if you iterate over it.
     808
     809        If flush is set, then the internal cache for attribute names will
    713810        be flushed. This may be necessary after the database schema or
    714811        the search path has been changed.
     
    735832            if not self._regtypes:
    736833                names = ((name, _simpletype(typ)) for name, typ in names)
    737             names = OrderedDict(names)
     834            names = AttrDict(names)
    738835            attnames[table] = names  # cache it
    739836        return names
     
    11921289        q = ' '.join(q)
    11931290        self._do_debug(q)
    1194         return self.query(q)
     1291        return self.db.query(q)
     1292
     1293    def get_as_list(self, table, what=None, where=None,
     1294            order=None, limit=None, offset=None, scalar=False):
     1295        """Get a table as a list.
     1296
     1297        This gets a convenient representation of the table as a list
     1298        of named tuples in Python.  You only need to pass the name of
     1299        the table (or any other SQL expression returning rows).  Note that
     1300        by default this will return the full content of the table which
     1301        can be huge and overflow your memory.  However, you can control
     1302        the amount of data returned using the other optional parameters.
     1303
     1304        The parameter 'what' can restrict the query to only return a
     1305        subset of the table columns.  It can be a string, list or a tuple.
     1306        The parameter 'where' can restrict the query to only return a
     1307        subset of the table rows.  It can be a string, list or a tuple
     1308        of SQL expressions that all need to be fulfilled.  The parameter
     1309        'order' specifies the ordering of the rows.  It can also be a
     1310        other string, list or a tuple.  If no ordering is specified,
     1311        the result will be ordered by the primary key(s) or all columns
     1312        if no primary key exists.  You can set 'order' to False if you
     1313        don't care about the ordering.  The parameters 'limit' and 'offset'
     1314        can be integers specifying the maximum number of rows returned
     1315        and a number of rows skipped over.
     1316
     1317        If you set the 'scalar' option to True, then instead of the
     1318        named tuples you will get the first items of these tuples.
     1319        This is useful if the result has only one column anyway.
     1320        """
     1321        if not table:
     1322            raise TypeError('The table name is missing')
     1323        if what:
     1324            if isinstance(what, (list, tuple)):
     1325                what = ', '.join(map(str, what))
     1326            if order is None:
     1327                order = what
     1328        else:
     1329            what = '*'
     1330        q = ['SELECT', what, 'FROM', table]
     1331        if where:
     1332            if isinstance(where, (list, tuple)):
     1333                where = ' AND '.join(map(str, where))
     1334            q.extend(['WHERE', where])
     1335        if order is None:
     1336            try:
     1337                order = self.pkey(table, True)
     1338            except (KeyError, ProgrammingError):
     1339                try:
     1340                    order = list(self.get_attnames(table))
     1341                except (KeyError, ProgrammingError):
     1342                    pass
     1343        if order:
     1344            if isinstance(order, (list, tuple)):
     1345                order = ', '.join(map(str, order))
     1346            q.extend(['ORDER BY', order])
     1347        if limit:
     1348            q.append('LIMIT %d' % limit)
     1349        if offset:
     1350            q.append('OFFSET %d' % offset)
     1351        q = ' '.join(q)
     1352        self._do_debug(q)
     1353        q = self.db.query(q)
     1354        res = q.namedresult()
     1355        if res and scalar:
     1356            res = [row[0] for row in res]
     1357        return res
     1358
     1359    def get_as_dict(self, table, keyname=None, what=None, where=None,
     1360            order=None, limit=None, offset=None, scalar=False):
     1361        """Get a table as a dictionary.
     1362
     1363        This method is similar to get_as_list(), but returns the table
     1364        as a Python dict instead of a Python list, which can be even
     1365        more convenient. The primary key column(s) of the table will
     1366        be used as the keys of the dictionary, while the other column(s)
     1367        will be the corresponding values.  The keys will be named tuples
     1368        if the table has a composite primary key.  The rows will be also
     1369        named tuples unless the 'scalar' option has been set to True.
     1370        With the optional parameter 'keyname' you can specify an alternative
     1371        set of columns to be used as the keys of the dictionary.  It must
     1372        be set as a string, list or a tuple.
     1373
     1374        If the Python version supports it, the dictionary will be an
     1375        OrderedDict using the order specified with the 'order' parameter
     1376        or the key column(s) if not specified.  You can set 'order' to False
     1377        if you don't care about the ordering.  In this case the returned
     1378        dictionary will be an ordinary one.
     1379        """
     1380        if not table:
     1381            raise TypeError('The table name is missing')
     1382        if not keyname:
     1383            try:
     1384                keyname = self.pkey(table, True)
     1385            except (KeyError, ProgrammingError):
     1386                raise _prg_error('Table %s has no primary key' % table)
     1387        if isinstance(keyname, basestring):
     1388            keyname = [keyname]
     1389        elif not isinstance(keyname, (list, tuple)):
     1390            raise KeyError('The keyname must be a string, list or tuple')
     1391        if what:
     1392            if isinstance(what, (list, tuple)):
     1393                what = ', '.join(map(str, what))
     1394            if order is None:
     1395                order = what
     1396        else:
     1397            what = '*'
     1398        q = ['SELECT', what, 'FROM', table]
     1399        if where:
     1400            if isinstance(where, (list, tuple)):
     1401                where = ' AND '.join(map(str, where))
     1402            q.extend(['WHERE', where])
     1403        if order is None:
     1404            order = keyname
     1405        if order:
     1406            if isinstance(order, (list, tuple)):
     1407                order = ', '.join(map(str, order))
     1408            q.extend(['ORDER BY', order])
     1409        if limit:
     1410            q.append('LIMIT %d' % limit)
     1411        if offset:
     1412            q.append('OFFSET %d' % offset)
     1413        q = ' '.join(q)
     1414        self._do_debug(q)
     1415        q = self.db.query(q)
     1416        res = q.getresult()
     1417        cls = OrderedDict if order else dict
     1418        if not res:
     1419            return cls()
     1420        keyset = set(keyname)
     1421        fields = q.listfields()
     1422        if not keyset.issubset(fields):
     1423            raise KeyError('Missing keyname in row')
     1424        keyind, rowind = [], []
     1425        for i, f in enumerate(fields):
     1426            (keyind if f in keyset else rowind).append(i)
     1427        keytuple = len(keyind) > 1
     1428        getkey = itemgetter(*keyind)
     1429        keys = map(getkey, res)
     1430        if scalar:
     1431            rowind = rowind[:1]
     1432            rowtuple = False
     1433        else:
     1434            rowtuple = len(rowind) > 1
     1435        if scalar or rowtuple:
     1436            getrow = itemgetter(*rowind)
     1437        else:
     1438            rowind = rowind[0]
     1439            getrow = lambda row: (row[rowind],)
     1440            rowtuple = True
     1441        rows = map(getrow, res)
     1442        if keytuple or rowtuple:
     1443            namedresult = get_namedresult()
     1444            if keytuple:
     1445                keys = namedresult(_MemoryQuery(keys, keyname))
     1446            if rowtuple:
     1447                fields = [f for f in fields if f not in keyset]
     1448                rows = namedresult(_MemoryQuery(rows, fields))
     1449        return cls(zip(keys, rows))
    11951450
    11961451    def notification_handler(self,
Note: See TracChangeset for help on using the changeset viewer.