diff --git a/src/pl/plpython/expected/plpython_spi.out b/src/pl/plpython/expected/plpython_spi.out index 671c24e33f..3cda9589b3 100644 --- a/src/pl/plpython/expected/plpython_spi.out +++ b/src/pl/plpython/expected/plpython_spi.out @@ -1,23 +1,3 @@ --- --- result objects --- -CREATE FUNCTION test_resultobject_access() RETURNS void -AS $$ -rv = plpy.execute("SELECT fname, lname, username FROM users ORDER BY username") -plpy.info([row for row in rv]) -rv[1] = dict([(k, v*2) for (k, v) in rv[1].items()]) -plpy.info([row for row in rv]) -$$ LANGUAGE plpythonu; -SELECT test_resultobject_access(); -INFO: [{'lname': 'doe', 'username': 'j_doe', 'fname': 'jane'}, {'lname': 'doe', 'username': 'johnd', 'fname': 'john'}, {'lname': 'smith', 'username': 'slash', 'fname': 'rick'}, {'lname': 'doe', 'username': 'w_doe', 'fname': 'willem'}] -CONTEXT: PL/Python function "test_resultobject_access" -INFO: [{'lname': 'doe', 'username': 'j_doe', 'fname': 'jane'}, {'lname': 'doedoe', 'username': 'johndjohnd', 'fname': 'johnjohn'}, {'lname': 'smith', 'username': 'slash', 'fname': 'rick'}, {'lname': 'doe', 'username': 'w_doe', 'fname': 'willem'}] -CONTEXT: PL/Python function "test_resultobject_access" - test_resultobject_access --------------------------- - -(1 row) - -- -- nested calls -- @@ -228,6 +208,61 @@ SELECT result_len_test($$UPDATE foo3 SET b= '' WHERE a = 2$$); 0 (1 row) +CREATE FUNCTION result_subscript_test() RETURNS void +AS $$ +result = plpy.execute("SELECT 1 AS c UNION SELECT 2 " + "UNION SELECT 3 UNION SELECT 4") + +plpy.info(result[1]['c']) +plpy.info(result[-1]['c']) + +plpy.info([item['c'] for item in result[1:3]]) +plpy.info([item['c'] for item in result[::2]]) + +result[-1] = {'c': 1000} +result[:2] = [{'c': 10}, {'c': 100}] +plpy.info([item['c'] for item in result[:]]) + +# raises TypeError, but the message differs on Python 2.6, so silence it +try: + plpy.info(result['foo']) +except TypeError: + pass +else: + assert False, "TypeError not raised" + +$$ LANGUAGE plpythonu; +SELECT result_subscript_test(); +INFO: 2 +CONTEXT: PL/Python function "result_subscript_test" +INFO: 4 +CONTEXT: PL/Python function "result_subscript_test" +INFO: [2, 3] +CONTEXT: PL/Python function "result_subscript_test" +INFO: [1, 3] +CONTEXT: PL/Python function "result_subscript_test" +INFO: [10, 100, 3, 1000] +CONTEXT: PL/Python function "result_subscript_test" + result_subscript_test +----------------------- + +(1 row) + +CREATE FUNCTION result_empty_test() RETURNS void +AS $$ +result = plpy.execute("select 1 where false") + +plpy.info(result[:]) + +$$ LANGUAGE plpythonu; +SELECT result_empty_test(); +INFO: [] +CONTEXT: PL/Python function "result_empty_test" + result_empty_test +------------------- + +(1 row) + -- cursor objects CREATE FUNCTION simple_cursor_test() RETURNS int AS $$ res = plpy.cursor("select fname, lname from users") diff --git a/src/pl/plpython/plpy_resultobject.c b/src/pl/plpython/plpy_resultobject.c index fcf8074228..deaddb7980 100644 --- a/src/pl/plpython/plpy_resultobject.c +++ b/src/pl/plpython/plpy_resultobject.c @@ -23,6 +23,8 @@ static PyObject *PLy_result_item(PyObject *arg, Py_ssize_t idx); static PyObject *PLy_result_slice(PyObject *arg, Py_ssize_t lidx, Py_ssize_t hidx); static int PLy_result_ass_item(PyObject *arg, Py_ssize_t idx, PyObject *item); static int PLy_result_ass_slice(PyObject *rg, Py_ssize_t lidx, Py_ssize_t hidx, PyObject *slice); +static PyObject *PLy_result_subscript(PyObject *arg, PyObject *item); +static int PLy_result_ass_subscript(PyObject* self, PyObject* item, PyObject* value); static char PLy_result_doc[] = { "Results of a PostgreSQL query" @@ -38,6 +40,12 @@ static PySequenceMethods PLy_result_as_sequence = { PLy_result_ass_slice, /* sq_ass_slice */ }; +static PyMappingMethods PLy_result_as_mapping = { + PLy_result_length, /* mp_length */ + PLy_result_subscript, /* mp_subscript */ + PLy_result_ass_subscript, /* mp_ass_subscript */ +}; + static PyMethodDef PLy_result_methods[] = { {"colnames", PLy_result_colnames, METH_NOARGS, NULL}, {"coltypes", PLy_result_coltypes, METH_NOARGS, NULL}, @@ -64,7 +72,7 @@ static PyTypeObject PLy_ResultType = { 0, /* tp_repr */ 0, /* tp_as_number */ &PLy_result_as_sequence, /* tp_as_sequence */ - 0, /* tp_as_mapping */ + &PLy_result_as_mapping, /* tp_as_mapping */ 0, /* tp_hash */ 0, /* tp_call */ 0, /* tp_str */ @@ -251,3 +259,19 @@ PLy_result_ass_slice(PyObject *arg, Py_ssize_t lidx, Py_ssize_t hidx, PyObject * rv = PyList_SetSlice(ob->rows, lidx, hidx, slice); return rv; } + +static PyObject * +PLy_result_subscript(PyObject *arg, PyObject *item) +{ + PLyResultObject *ob = (PLyResultObject *) arg; + + return PyObject_GetItem(ob->rows, item); +} + +static int +PLy_result_ass_subscript(PyObject *arg, PyObject *item, PyObject *value) +{ + PLyResultObject *ob = (PLyResultObject *) arg; + + return PyObject_SetItem(ob->rows, item, value); +} diff --git a/src/pl/plpython/sql/plpython_spi.sql b/src/pl/plpython/sql/plpython_spi.sql index 7be2fbff71..6250e90d19 100644 --- a/src/pl/plpython/sql/plpython_spi.sql +++ b/src/pl/plpython/sql/plpython_spi.sql @@ -1,18 +1,3 @@ --- --- result objects --- - -CREATE FUNCTION test_resultobject_access() RETURNS void -AS $$ -rv = plpy.execute("SELECT fname, lname, username FROM users ORDER BY username") -plpy.info([row for row in rv]) -rv[1] = dict([(k, v*2) for (k, v) in rv[1].items()]) -plpy.info([row for row in rv]) -$$ LANGUAGE plpythonu; - -SELECT test_resultobject_access(); - - -- -- nested calls -- @@ -147,6 +132,42 @@ SELECT result_len_test($$CREATE TEMPORARY TABLE foo3 (a int, b text)$$); SELECT result_len_test($$INSERT INTO foo3 VALUES (1, 'one'), (2, 'two')$$); SELECT result_len_test($$UPDATE foo3 SET b= '' WHERE a = 2$$); +CREATE FUNCTION result_subscript_test() RETURNS void +AS $$ +result = plpy.execute("SELECT 1 AS c UNION SELECT 2 " + "UNION SELECT 3 UNION SELECT 4") + +plpy.info(result[1]['c']) +plpy.info(result[-1]['c']) + +plpy.info([item['c'] for item in result[1:3]]) +plpy.info([item['c'] for item in result[::2]]) + +result[-1] = {'c': 1000} +result[:2] = [{'c': 10}, {'c': 100}] +plpy.info([item['c'] for item in result[:]]) + +# raises TypeError, but the message differs on Python 2.6, so silence it +try: + plpy.info(result['foo']) +except TypeError: + pass +else: + assert False, "TypeError not raised" + +$$ LANGUAGE plpythonu; + +SELECT result_subscript_test(); + +CREATE FUNCTION result_empty_test() RETURNS void +AS $$ +result = plpy.execute("select 1 where false") + +plpy.info(result[:]) + +$$ LANGUAGE plpythonu; + +SELECT result_empty_test(); -- cursor objects