added DictCursor support. fixed a previously unknown bug in executemany(). fixed bug in python3 version of cursor.description

petehunt · petehunt · commit b3b28263ef84 · 2011-03-01T18:42:47.000Z
diff --git a/pymysql/connections.py b/pymysql/connections.py
@@ -189,7 +189,8 @@ class MysqlPacket(object):
   from the network socket, removes packet header and provides an interface
   for reading/parsing the packet results."""
 
-  def __init__(self, socket):
+  def __init__(self, socket, connection):
+    self.connection = connection
     self.__position = 0
     self.__recv_packet(socket)
     del socket
@@ -354,7 +355,7 @@ def __parse_field_descriptor(self):
     self.db = self.read_length_coded_string()
     self.table_name = self.read_length_coded_string()
     self.org_table = self.read_length_coded_string()
-    self.name = self.read_length_coded_string()
+    self.name = self.read_length_coded_string().decode(self.connection.charset)
     self.org_name = self.read_length_coded_string()
     self.advance(1)  # non-null filler
     self.charsetnr = struct.unpack('<H', self.read(2))[0]
@@ -578,8 +579,10 @@ def literal(self, obj):
         ''' Alias for escape() '''
         return escape_item(obj, self.charset)
 
-    def cursor(self):
+    def cursor(self, cursor=None):
         ''' Create a new cursor to execute queries with '''
+        if cursor:
+            return cursor(self)
         return self.cursorclass(self)
 
     def __enter__(self):
@@ -676,7 +679,7 @@ def read_packet(self, packet_type=MysqlPacket):
       #       socket.recv(large_number)?  if so, maybe we should buffer
       #       the socket.recv() (though that obviously makes memory management
       #       more complicated.
-      packet = packet_type(self.socket)
+      packet = packet_type(self.socket, self)
       packet.check_error()
       return packet
 
@@ -762,7 +765,7 @@ def _send_authentication(self):
 
         sock.send(data)
 
-        auth_packet = MysqlPacket(sock)
+        auth_packet = MysqlPacket(sock, self)
         auth_packet.check_error()
         if DEBUG: auth_packet.dump()
 
@@ -777,7 +780,7 @@ def _send_authentication(self):
             data = pack_int24(len(data)) + int2byte(next_packet) + data
 
             sock.send(data)
-            auth_packet = MysqlPacket(sock)
+            auth_packet = MysqlPacket(sock, self)
             auth_packet.check_error()
             if DEBUG: auth_packet.dump()
 
@@ -798,7 +801,7 @@ def get_proto_info(self):
     def _get_server_information(self):
         sock = self.socket
         i = 0
-        packet = MysqlPacket(sock)
+        packet = MysqlPacket(sock, self)
         data = packet.get_all_data()
 
         if DEBUG: dump_packet(data)
diff --git a/pymysql/cursors.py b/pymysql/cursors.py
@@ -113,12 +113,12 @@ def execute(self, query, args=None):
     def executemany(self, query, args):
         ''' Run several data against one query '''
         del self.messages[:]
-        conn = self._get_db()
+        #conn = self._get_db()
         if not args:
             return
-        charset = conn.charset
-        if isinstance(query, unicode):
-            query = query.encode(charset)
+        #charset = conn.charset
+        #if isinstance(query, unicode):
+        #    query = query.encode(charset)
 
         self.rowcount = sum([ self.execute(query, arg) for arg in args ])
         return self.rowcount
@@ -248,3 +248,44 @@ def __iter__(self):
     InternalError = InternalError
     ProgrammingError = ProgrammingError
     NotSupportedError = NotSupportedError
+
+class DictCursor(Cursor):
+    """A cursor which returns results as a dictionary"""
+
+    def execute(self, query, args=None):
+        result = super(DictCursor, self).execute(query, args)
+        if self.description:
+            self._fields = [ field[0] for field in self.description ]
+        return result
+
+    def fetchone(self):
+        ''' Fetch the next row '''
+        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]))
+        self.rownumber += 1
+        return result
+
+    def fetchmany(self, size=None):
+        ''' Fetch several rows '''
+        self._check_executed()
+        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] ]
+        self.rownumber = min(end, len(self._rows))
+        return tuple(result)
+
+    def fetchall(self):
+        ''' Fetch all the rows '''
+        self._check_executed()
+        if self._rows is None:
+            return None
+        if self.rownumber:
+            result = [ dict(zip(self._fields, r)) for r in self._rows[self.rownumber:] ]
+        else:
+            result = [ dict(zip(self._fields, r)) for r in self._rows ]
+        self.rownumber = len(self._rows)
+        return tuple(result)
+
diff --git a/pymysql/tests/__init__.py b/pymysql/tests/__init__.py
@@ -1,6 +1,7 @@
 from pymysql.tests.test_issues import *
 from pymysql.tests.test_example import *
 from pymysql.tests.test_basic import *
+from pymysql.tests.test_DictCursor import *
 
 if __name__ == "__main__":
     import unittest
diff --git a/pymysql/tests/test_DictCursor.py b/pymysql/tests/test_DictCursor.py
@@ -0,0 +1,48 @@
+from pymysql.tests import base
+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]
+        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, 02, 6, 23, 04, 56)}
+        jim =  {'name':'jim','age':56,'DOB':datetime.datetime(1955, 05, 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")
+            # get all 3 row via fetchall
+            c.execute("SELECT * from dictcursor")
+            r = c.fetchall()
+            self.assertEqual((bob,jim,fred), r, "fetchall 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")
+
+__all__ = ["TestDictCursor"]
+
+if __name__ == "__main__":
+    import unittest
+    unittest.main()
