Merge pull request #158 from methane/dict-type

Add DictCursor.dict_type to override it.
This commit is contained in:
INADA Naoki
2013-09-22 18:49:40 -07:00
2 changed files with 97 additions and 54 deletions

View File

@@ -12,9 +12,10 @@ if PY2:
else:
import io
from .err import Warning, Error, InterfaceError, DataError, \
DatabaseError, OperationalError, IntegrityError, InternalError, \
NotSupportedError, ProgrammingError
from .err import (
Warning, Error, InterfaceError, DataError,
DatabaseError, OperationalError, IntegrityError, InternalError,
NotSupportedError, ProgrammingError)
insert_values = re.compile(r'\svalues\s*(\(.+\))', re.IGNORECASE)
@@ -115,9 +116,8 @@ class Cursor(object):
result = 0
try:
result = self._query(query)
except:
exc, value, tb = exc_info()
del tb
except Exception:
exc, value = exc_info()[:2]
self.errorhandler(self, exc, value)
self._executed = query
@@ -261,6 +261,8 @@ class Cursor(object):
class DictCursor(Cursor):
"""A cursor which returns results as a dictionary"""
# You can override this to use OrderedDict or other dict-like types.
dict_type = dict
def execute(self, query, args=None):
result = super(DictCursor, self).execute(query, args)
@@ -273,7 +275,7 @@ class DictCursor(Cursor):
self._check_executed()
if self._rows is None or self.rownumber >= len(self._rows):
return None
result = dict(zip(self._fields, self._rows[self.rownumber]))
result = self.dict_type(zip(self._fields, self._rows[self.rownumber]))
self.rownumber += 1
return result
@@ -283,9 +285,9 @@ class DictCursor(Cursor):
if self._rows is None:
return None
end = self.rownumber + (size or self.arraysize)
result = [ dict(zip(self._fields, r)) for r in self._rows[self.rownumber:end] ]
result = [self.dict_type(zip(self._fields, r)) for r in self._rows[self.rownumber:end]]
self.rownumber = min(end, len(self._rows))
return tuple(result)
return result
def fetchall(self):
''' Fetch all the rows '''
@@ -293,11 +295,11 @@ class DictCursor(Cursor):
if self._rows is None:
return None
if self.rownumber:
result = [ dict(zip(self._fields, r)) for r in self._rows[self.rownumber:] ]
result = [self.dict_type(zip(self._fields, r)) for r in self._rows[self.rownumber:]]
else:
result = [ dict(zip(self._fields, r)) for r in self._rows ]
result = [self.dict_type(zip(self._fields, r)) for r in self._rows]
self.rownumber = len(self._rows)
return tuple(result)
return result
class SSCursor(Cursor):
"""

View File

@@ -3,53 +3,94 @@ import pymysql.cursors
import datetime
class TestDictCursor(base.PyMySQLTestCase):
def test_DictCursor(self):
#all assert test compare to the structure as would come out from MySQLdb
conn = self.connections[0]
class TestDictCursor(base.PyMySQLTestCase):
bob = {'name': 'bob', 'age': 21, 'DOB': datetime.datetime(1990, 2, 6, 23, 4, 56)}
jim = {'name': 'jim', 'age': 56, 'DOB': datetime.datetime(1955, 5, 9, 13, 12, 45)}
fred = {'name': 'fred', 'age': 100, 'DOB': datetime.datetime(1911, 9, 12, 1, 1, 1)}
def setUp(self):
super(TestDictCursor, self).setUp()
self.conn = conn = self.connections[0]
c = conn.cursor(pymysql.cursors.DictCursor)
# create a table ane some data to query
c.execute("""CREATE TABLE dictcursor (name char(20), age int , DOB datetime)""")
data = (("bob",21,"1990-02-06 23:04:56"),
("jim",56,"1955-05-09 13:12:45"),
("fred",100,"1911-09-12 01:01:01"))
bob = {'name':'bob','age':21,'DOB':datetime.datetime(1990, 2, 6, 23, 4, 56)}
jim = {'name':'jim','age':56,'DOB':datetime.datetime(1955, 5, 9, 13, 12, 45)}
fred = {'name':'fred','age':100,'DOB':datetime.datetime(1911, 9, 12, 1, 1, 1)}
try:
c.executemany("insert into dictcursor values (%s,%s,%s)", data)
# try an update which should return no rows
c.execute("update dictcursor set age=20 where name='bob'")
bob['age'] = 20
# pull back the single row dict for bob and check
c.execute("SELECT * from dictcursor where name='bob'")
r = c.fetchone()
self.assertEqual(bob,r,"fetchone via DictCursor failed")
# same again, but via fetchall => tuple)
c.execute("SELECT * from dictcursor where name='bob'")
r = c.fetchall()
self.assertEqual((bob,),r,"fetch a 1 row result via fetchall failed via DictCursor")
# same test again but iterate over the
c.execute("SELECT * from dictcursor where name='bob'")
for r in c:
self.assertEqual(bob, r,"fetch a 1 row result via iteration failed via DictCursor")
# get all 3 row via fetchall
c.execute("SELECT * from dictcursor")
r = c.fetchall()
self.assertEqual((bob,jim,fred), r, "fetchall failed via DictCursor")
#same test again but do a list comprehension
c.execute("SELECT * from dictcursor")
r = [x for x in c]
self.assertEqual([bob,jim,fred], r, "list comprehension failed via DictCursor")
# get all 2 row via fetchmany
c.execute("SELECT * from dictcursor")
r = c.fetchmany(2)
self.assertEqual((bob,jim), r, "fetchmany failed via DictCursor")
finally:
c.execute("drop table dictcursor")
data = [("bob", 21, "1990-02-06 23:04:56"),
("jim", 56, "1955-05-09 13:12:45"),
("fred", 100, "1911-09-12 01:01:01")]
c.executemany("insert into dictcursor values (%s,%s,%s)", data)
def tearDown(self):
c = self.conn.cursor()
c.execute("drop table dictcursor")
super(TestDictCursor, self).tearDown()
def test_DictCursor(self):
bob, jim, fred = self.bob.copy(), self.jim.copy(), self.fred.copy()
#all assert test compare to the structure as would come out from MySQLdb
conn = self.conn
c = conn.cursor(pymysql.cursors.DictCursor)
# try an update which should return no rows
c.execute("update dictcursor set age=20 where name='bob'")
bob['age'] = 20
# pull back the single row dict for bob and check
c.execute("SELECT * from dictcursor where name='bob'")
r = c.fetchone()
self.assertEqual(bob, r, "fetchone via DictCursor failed")
# same again, but via fetchall => tuple)
c.execute("SELECT * from dictcursor where name='bob'")
r = c.fetchall()
self.assertEqual([bob], r, "fetch a 1 row result via fetchall failed via DictCursor")
# same test again but iterate over the
c.execute("SELECT * from dictcursor where name='bob'")
for r in c:
self.assertEqual(bob, r, "fetch a 1 row result via iteration failed via DictCursor")
# get all 3 row via fetchall
c.execute("SELECT * from dictcursor")
r = c.fetchall()
self.assertEqual([bob,jim,fred], r, "fetchall failed via DictCursor")
#same test again but do a list comprehension
c.execute("SELECT * from dictcursor")
r = list(c)
self.assertEqual([bob,jim,fred], r, "DictCursor should be iterable")
# get all 2 row via fetchmany
c.execute("SELECT * from dictcursor")
r = c.fetchmany(2)
self.assertEqual([bob, jim], r, "fetchmany failed via DictCursor")
def test_custom_dict(self):
class MyDict(dict): pass
class MyDictCursor(pymysql.cursors.DictCursor):
dict_type = MyDict
keys = ['name', 'age', 'DOB']
bob = MyDict([(k, self.bob[k]) for k in keys])
jim = MyDict([(k, self.jim[k]) for k in keys])
fred = MyDict([(k, self.fred[k]) for k in keys])
cur = self.conn.cursor(MyDictCursor)
cur.execute("SELECT * FROM dictcursor WHERE name='bob'")
r = cur.fetchone()
self.assertEqual(bob, r, "fetchone() returns MyDictCursor")
cur.execute("SELECT * FROM dictcursor")
r = cur.fetchall()
self.assertEqual([bob, jim, fred], r,
"fetchall failed via MyDictCursor")
cur.execute("SELECT * FROM dictcursor")
r = list(cur)
self.assertEqual([bob, jim, fred], r,
"list failed via MyDictCursor")
cur.execute("SELECT * FROM dictcursor")
r = cur.fetchmany(2)
self.assertEqual([bob, jim], r,
"list failed via MyDictCursor")
__all__ = ["TestDictCursor"]
if __name__ == "__main__":
import unittest