Skip to content

Commit c29b04f

Browse files
mkleehammermkleehammer
committed
Cursor.executemany now accepts an iterator or a generator. closes mkleehammer#12 
1 parent 6463a91 commit c29b04f

3 files changed

Lines changed: 91 additions & 29 deletions

File tree

‎src/cursor.cpp‎

Lines changed: 54 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
#include "dbspecific.h"
2626
#include "sqlwchar.h"
2727
#include <datetime.h>
28+
#include "wrapper.h"
2829

2930
enum
3031
{
@@ -959,32 +960,66 @@ static PyObject* Cursor_executemany(PyObject* self, PyObject* args)
959960
return 0;
960961
}
961962

962-
if (!IsSequence(param_seq))
963+
if (IsSequence(param_seq))
963964
{
964-
PyErr_SetString(ProgrammingError, "The second parameter to executemany must be a sequence.");
965-
return 0;
966-
}
965+
Py_ssize_t c = PySequence_Size(param_seq);
967966

968-
Py_ssize_t c = PySequence_Size(param_seq);
967+
if (c == 0)
968+
{
969+
PyErr_SetString(ProgrammingError, "The second parameter to executemany must not be empty.");
970+
return 0;
971+
}
969972

970-
if (c == 0)
971-
{
972-
PyErr_SetString(ProgrammingError, "The second parameter to executemany must not be empty.");
973-
return 0;
973+
for (Py_ssize_t i = 0; i < c; i++)
974+
{
975+
PyObject* params = PySequence_GetItem(param_seq, i);
976+
PyObject* result = execute(cursor, pSql, params, false);
977+
bool success = result != 0;
978+
Py_XDECREF(result);
979+
Py_DECREF(params);
980+
if (!success)
981+
{
982+
cursor->rowcount = -1;
983+
return 0;
984+
}
985+
}
974986
}
975-
976-
for (Py_ssize_t i = 0; i < c; i++)
987+
else if (PyGen_Check(param_seq) || PyIter_Check(param_seq))
977988
{
978-
PyObject* params = PySequence_GetItem(param_seq, i);
979-
PyObject* result = execute(cursor, pSql, params, false);
980-
bool success = result != 0;
981-
Py_XDECREF(result);
982-
Py_DECREF(params);
983-
if (!success)
989+
Object iter;
990+
991+
if (PyGen_Check(param_seq))
984992
{
985-
cursor->rowcount = -1;
986-
return 0;
993+
iter = PyObject_GetIter(param_seq);
994+
}
995+
else
996+
{
997+
iter = param_seq;
998+
Py_INCREF(param_seq);
999+
}
1000+
1001+
Object params;
1002+
1003+
while (params.Attach(PyIter_Next(iter)))
1004+
{
1005+
PyObject* result = execute(cursor, pSql, params, false);
1006+
bool success = result != 0;
1007+
Py_XDECREF(result);
1008+
1009+
if (!success)
1010+
{
1011+
cursor->rowcount = -1;
1012+
return 0;
1013+
}
9871014
}
1015+
1016+
if (PyErr_Occurred())
1017+
return 0;
1018+
}
1019+
else
1020+
{
1021+
PyErr_SetString(ProgrammingError, "The second parameter to executemany must be a sequence, iterator, or generator.");
1022+
return 0;
9881023
}
9891024

9901025
cursor->rowcount = -1;

‎src/wrapper.h‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,13 @@ class Object
3131

3232
bool IsValid() const { return p != 0; }
3333

34-
void Attach(PyObject* _p)
34+
bool Attach(PyObject* _p)
3535
{
36+
// Returns true if the new pointer is non-zero.
37+
3638
Py_XDECREF(p);
3739
p = _p;
40+
return (_p != 0);
3841
}
3942

4043
PyObject* Detach()

‎tests2/pgtests.py‎

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ def setUp(self):
5757
self.cnxn.commit()
5858
except:
5959
pass
60-
60+
6161
self.cnxn.rollback()
6262

6363

@@ -192,7 +192,7 @@ def test_negative_decimal_scale(self):
192192

193193
def _exec(self):
194194
self.cursor.execute(self.sql)
195-
195+
196196
def test_close_cnxn(self):
197197
"""Make sure using a Cursor after closing its connection doesn't crash."""
198198

@@ -201,7 +201,7 @@ def test_close_cnxn(self):
201201
self.cursor.execute("select * from t1")
202202

203203
self.cnxn.close()
204-
204+
205205
# Now that the connection is closed, we expect an exception. (If the code attempts to use
206206
# the HSTMT, we'll get an access violation instead.)
207207
self.sql = "select * from t1"
@@ -263,13 +263,13 @@ def test_rowcount_select(self):
263263
# PostgreSQL driver fails here?
264264
# def test_rowcount_reset(self):
265265
# "Ensure rowcount is reset to -1"
266-
#
266+
#
267267
# self.cursor.execute("create table t1(i int)")
268268
# count = 4
269269
# for i in range(count):
270270
# self.cursor.execute("insert into t1 values (?)", i)
271271
# self.assertEquals(self.cursor.rowcount, 1)
272-
#
272+
#
273273
# self.cursor.execute("create table t2(i int)")
274274
# self.assertEquals(self.cursor.rowcount, -1)
275275

@@ -291,7 +291,7 @@ def test_lower_case(self):
291291

292292
# Put it back so other tests don't fail.
293293
pyodbc.lowercase = False
294-
294+
295295
def test_row_description(self):
296296
"""
297297
Ensure Cursor.description is accessible as Row.cursor_description.
@@ -303,7 +303,7 @@ def test_row_description(self):
303303

304304
row = self.cursor.execute("select * from t1").fetchone()
305305
self.assertEquals(self.cursor.description, row.cursor_description)
306-
306+
307307

308308
def test_executemany(self):
309309
self.cursor.execute("create table t1(a int, b varchar(10))")
@@ -336,10 +336,34 @@ def test_executemany_failure(self):
336336
params = [ (1, 'good'),
337337
('error', 'not an int'),
338338
(3, 'good') ]
339-
339+
340340
self.failUnlessRaises(pyodbc.Error, self.cursor.executemany, "insert into t1(a, b) value (?, ?)", params)
341341

342-
342+
343+
def test_executemany_generator(self):
344+
self.cursor.execute("create table t1(a int)")
345+
346+
self.cursor.executemany("insert into t1(a) values (?)", ((i,) for i in range(4)))
347+
348+
row = self.cursor.execute("select min(a) mina, max(a) maxa from t1").fetchone()
349+
350+
self.assertEqual(row.mina, 0)
351+
self.assertEqual(row.maxa, 3)
352+
353+
354+
def test_executemany_iterator(self):
355+
self.cursor.execute("create table t1(a int)")
356+
357+
values = [ (i,) for i in range(4) ]
358+
359+
self.cursor.executemany("insert into t1(a) values (?)", iter(values))
360+
361+
row = self.cursor.execute("select min(a) mina, max(a) maxa from t1").fetchone()
362+
363+
self.assertEqual(row.mina, 0)
364+
self.assertEqual(row.maxa, 3)
365+
366+
343367
def test_row_slicing(self):
344368
self.cursor.execute("create table t1(a int, b int, c int, d int)");
345369
self.cursor.execute("insert into t1 values(1,2,3,4)")

0 commit comments

Comments
 (0)