Package rdkit :: Package Dbase :: Module DbResultSet

[hide private]
[frames] | no frames]

# Source Code for Module rdkit.Dbase.DbResultSet

      1  #
      2  #  Copyright (C) 2003  Greg Landrum and Rational Discovery LLC
      3  #
      4  """ defines class _DbResultSet_ for lazy interactions with Db query results
      5
      6  **Note**
      7
      8    this uses the Python iterator interface, so you'll need python 2.2 or above.
      9
     10  """
     11  from __future__ import print_function
     12  import sys
     13  from rdkit.Dbase import DbInfo
     14
     15 -class DbResultBase(object):
     16 -  def __init__(self,cursor,conn,cmd,removeDups=-1,transform=None,extras=None):
     17      self.cursor = cursor
     18      self.removeDups = removeDups
     19      self.transform = transform
     20      self.cmd = cmd
     21      self.conn=conn
     22      self.extras = extras
     23      self.Reset()
     24      self._initColumnNamesAndTypes()
     25      self.Reset()
     26
     27 -  def Reset(self):
     28      """ implement in subclasses
     29
     30      """
     31      try:
     32        if not self.extras:
     33          self.cursor.execute(self.cmd)
     34        else:
     35          self.cursor.execute(self.cmd,self.extras)
     36      except Exception:
     37        sys.stderr.write('the command "%s" generated errors:\n'%(self.cmd))
     38        import traceback
     39        traceback.print_exc()
     40
     41
     42 -  def __iter__(self):
     43      self.Reset()
     44      return self
     45
     46 -  def _initColumnNamesAndTypes(self):
     47      self.colNames = []
     48      self.colTypes = []
     49      for cName,cType in DbInfo.GetColumnInfoFromCursor(self.cursor):
     50        self.colNames.append(cName)
     51        self.colTypes.append(cType)
     52      self.colNames = tuple(self.colNames)
     53      self.colTypes = tuple(self.colTypes)
     54
     55 -  def GetColumnNames(self):
     56      return self.colNames
     57 -  def GetColumnTypes(self):
     58      return self.colTypes
     59 -  def GetColumnNamesAndTypes(self):
     60      res = [None]*len(self.colNames)
     61      for i in range(len(self.colNames)):
     62        res[i] = self.colNames[i],self.colTypes[i]
     63      return tuple(res)
     64
     65
     66 -class DbResultSet(DbResultBase):
     67    """ Only supports forward iteration
     68
     69    """
     70 -  def __init__(self,*args,**kwargs):
     71      DbResultBase.__init__(self,*args,**kwargs)
     72      self.seen = []
     73      self._stopped = 0
     74 -  def Reset(self):
     75      self._stopped = 0
     76      DbResultBase.Reset(self)
     77
     78 -  def next(self):
     79      if self._stopped:
     80        raise StopIteration
     81      r = None
     82      while r is None:
     83        r = self.cursor.fetchone()
     84        if not r:
     85          self._stopped = 1
     86          raise StopIteration
     87        if self.transform is not None:
     88          r = self.transform(r)
     89        if self.removeDups>=0:
     90          v = r[self.removeDups]
     91          if v in self.seen:
     92            r = None
     93          else:
     94            self.seen.append(v)
     95      return r
     96
     97    __next__ = next # PY3
     98
     99
    100 -class RandomAccessDbResultSet(DbResultBase):
    101    """ Supports random access
    102
    103    """
    104 -  def __init__(self,*args,**kwargs):
    105      DbResultBase.__init__(self,*args,**kwargs)
    106      self.results = []
    107      self.seen = []
    108      self._pos = -1
    109
    110 -  def Reset(self):
    111      self._pos = -1
    112      if self.cursor is not None:
    113        DbResultBase.Reset(self)
    114
    115 -  def _finish(self):
    116      if self.cursor:
    117        #sys.stderr.write('_finish:\n')
    118        r = self.cursor.fetchone()
    119        while r:
    120          if self.transform is not None:
    121            r =  self.transform(r)
    122          if self.removeDups >=0:
    123            v = r[self.removeDups]
    124            if v not in self.seen:
    125              self.seen.append(v)
    126              self.results.append(r)
    127          else:
    128            self.results.append(r)
    129          r = self.cursor.fetchone()
    130        self.cursor = None
    131 -  def __getitem__(self,idx):
    132      if idx < 0: raise IndexError("negative indices not supported")
    133      if self.cursor is None:
    134        if len(self.results):
    135          if idx >= len(self.results):
    136            raise IndexError('index %d too large (%d max)'%(idx,len(self.results)))
    137        else:
    138          raise ValueError('Invalid cursor')
    139
    140      while idx >= len(self.results):
    141        r = None
    142        while r is None:
    143          r = self.cursor.fetchone()
    144          if not r:
    145            self.cursor = None
    146            raise IndexError('index %d too large (%d max)'%(idx,len(self.results)))
    147
    148          if self.transform is not None:
    149            r =  self.transform(r)
    150          if self.removeDups>=0:
    151            v = r[self.removeDups]
    152            if v in self.seen:
    153              r = None
    154            else:
    155              self.results.append(r)
    156              self.seen.append(v)
    157          else:
    158            self.results.append(r)
    159
    160      return self.results[idx]
    161
    162 -  def __len__(self):
    163      if self.results is None:
    164        raise ValueError("len() not supported for noMemory Results Sets")
    165      self._finish()
    166      return len(self.results)
    167
    168 -  def next(self):
    169      self._pos += 1
    170      res = None
    171      if self._pos < len(self):
    172        res = self.results[self._pos]
    173      else:
    174        raise StopIteration
    175      return res
    176
    177    __next__ = next # PY3
    178
    179
    180  if __name__ == '__main__':
    181    from rdkit.Dbase.DbConnection import DbConnect
    182    conn = DbConnect('TEST.GDB')
    183    curs = conn.GetCursor()
    184    print('curs:',repr(curs))
    185    curs.execute('select * from ten_elements')
    186    set = RandomAccessDbResultSet(curs)
    187    for i in range(12):
    188      try:
    189        val = set[i]
    190      except IndexError:
    191        assert i >= 10
    192
    193    print('use len')
    194    curs = conn.GetCursor()
    195    curs.execute('select * from ten_elements')
    196    set = RandomAccessDbResultSet(curs)
    197    for i in range(len(set)):
    198      val = set[i]
    199
    200    print('use iter')
    201    curs = conn.GetCursor()
    202    curs.execute('select * from ten_elements')
    203    set = DbResultSet(curs)
    204    for thing in set:
    205      id,val = thing
    206
    207    print('dups')
    208    curs = conn.GetCursor()
    209    curs.execute('select * from ten_elements_dups')
    210    set = DbResultSet(curs)
    211    r = []
    212    for thing in set:
    213      r.append(thing)
    214    assert len(r)==20
    215
    216    curs = conn.GetCursor()
    217    curs.execute('select * from ten_elements_dups')
    218    set = DbResultSet(curs,removeDups=0)
    219    r = []
    220    for thing in set:
    221      r.append(thing)
    222    assert len(r)==10
    223
    224    curs = conn.GetCursor()
    225    curs.execute('select * from ten_elements_dups')
    226    set = RandomAccessDbResultSet(curs,removeDups=0)
    227    assert len(set)==10
    228    assert set[0] == (0,11)
    229
    230    curs = conn.GetCursor()
    231    curs.execute('select * from ten_elements_dups')
    232    set = RandomAccessDbResultSet(curs,removeDups=0)
    233    assert set[0] == (0,11)
    234    assert set[1] == (2,21)
    235    assert set[5] == (10,61)
    236
    237    curs = conn.GetCursor()
    238    curs.execute('select * from ten_elements_dups')
    239    set = RandomAccessDbResultSet(curs)
    240    assert set[0] == (0,11)
    241    assert set[1] == (0,11)
    242