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: 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