fix: improve set_result() proposal:

- make it async on async cursors (consistently with `results()`,
  inconsistently with `nextset()`, but the latter being sync on async
  cursors was a design mistake);
- raise IndexError, consistently with sequences, better error message;
- fix test, which were simply broken and the OP didn't bother to fix
  them;
- documentation improved.

diff --git a/docs/api/cursors.rst b/docs/api/cursors.rst
index 6abb2b86c7315c4a5a7cc711ab3087129df4c2ed..546a4fcd097c2c7e98f0e412c3394bd8338fbb73 100644 (file)
--- a/docs/api/cursors.rst
+++ b/docs/api/cursors.rst
@@ -264,14 +264,6 @@ The `!Cursor` class
     .. automethod:: fetchmany
     .. automethod:: fetchall
     .. automethod:: nextset
-    .. automethod:: set_result
-
-        Move to a specific result set if `execute()` returned multiple result
-        sets. The parameter `result_no` specifies the zero-based index of the
-        desired result set. Negative values are supported and refer to result
-        sets counted from the end.
-
-        .. versionadded:: 3.3
 
     .. automethod:: results
 
@@ -296,6 +288,10 @@ The `!Cursor` class
             In previous version you may call `nextset()` in a loop until it
             returns a false value.
 
+    .. automethod:: set_result
+
+        .. versionadded:: 3.3
+
     .. automethod:: scroll
 
     .. attribute:: pgresult
@@ -571,6 +567,7 @@ semantic with an `!async` interface. The main interface is described in
     .. automethod:: fetchmany
     .. automethod:: fetchall
     .. automethod:: results
+    .. automethod:: set_result
     .. automethod:: scroll
 
     .. note::

diff --git a/docs/news.rst b/docs/news.rst
index 0e891ac8e358783976635bfd87062e4b31181188..cfdc1c4476767d175ffd31d77a2b35d7a50e4b3c 100644 (file)
--- a/docs/news.rst
+++ b/docs/news.rst
@@ -19,11 +19,11 @@ Psycopg 3.3.0 (unreleased)
 - More flexible :ref:`composite adaptation<adapt-composite>`: it is now possible
   to adapt Python objects to PostgreSQL composites and back even if they are not
   sequences or if they take keyword arguments (:ticket:`#932`, :ticket:`#1202`).
-- Cursors are now iterators, not only iterables. This means you can call
+- Cursors are now *iterators*, not only *iterables*. This means you can call
   ``next(cur)`` to fetch the next row (:ticket:`#1064`).
-- Add `Cursor.results()` to iterate over the result sets of the queries
-  executed though `~Cursor.executemany()` or `~Cursor.execute()`
-  (:ticket:`#1080`).
+- Add `Cursor.set_result()` and `Cursor.results()` to move across the result
+  sets of queries executed though `~Cursor.executemany()` or
+  `~Cursor.execute()` with multiple statements (:tickets:`#1080, #1170`).
 
 .. rubric:: New libpq wrapper features
 

diff --git a/psycopg/psycopg/_cursor_base.py b/psycopg/psycopg/_cursor_base.py
index 5b76744dd15240dbf771e59a1db714cfa49dd338..31bb751220d93979c9881dab8d4312d1d71c1c39 100644 (file)
--- a/psycopg/psycopg/_cursor_base.py
+++ b/psycopg/psycopg/_cursor_base.py
@@ -165,21 +165,6 @@ class BaseCursor(Generic[ConnectionType, Row]):
         else:
             return None
 
-    def set_result(self, result_no: int) -> None:
-        """
-        Move to a specific result set if `execute()` returned multiple result sets.
-
-        Args:
-            result_no (int): Zero based index, supports negative values
-        """
-        total_results = len(self._results)
-        if result_no < 0:
-            result_no = total_results + result_no
-        if 0 <= result_no < total_results:
-            self._select_current_result(result_no)
-        else:
-            raise ValueError("result_no value not in available results range")
-
     @property
     def statusmessage(self) -> str | None:
         """

diff --git a/psycopg/psycopg/cursor.py b/psycopg/psycopg/cursor.py
index a3383dbaecdae64b4c2ecd293c7a512192abafeb..255075087154bae986706afdf10e06659d2cebc6 100644 (file)
--- a/psycopg/psycopg/cursor.py
+++ b/psycopg/psycopg/cursor.py
@@ -210,6 +210,32 @@ class Cursor(BaseCursor["Connection[Any]", Row]):
                 if not self.nextset():
                     break
 
+    def set_result(self, index: int) -> Self:
+        """
+        Move to a specific result set.
+
+        :arg index: index of the result to go to
+        :type index: `!int`
+
+        More than one result will be available after executing calling
+        `executemany()` or `execute()` with more than one query.
+
+        `!index` is 0-based and supports negative values, counting from the end,
+        the same way you can index items in a list.
+
+        The function returns self, so that the result may be followed by a
+        fetch operation. See `results()` for details.
+        """
+        if not -len(self._results) <= index < len(self._results):
+            raise IndexError(
+                f"index {index} out of range: {len(self._results)} result(s) available"
+            )
+        if index < 0:
+            index = len(self._results) + index
+
+        self._select_current_result(index)
+        return self
+
     def fetchone(self) -> Row | None:
         """
         Return the next record from the current result set.

diff --git a/psycopg/psycopg/cursor_async.py b/psycopg/psycopg/cursor_async.py
index 22dd9458321a097faf72810f951b39297d3eff82..c7bca86f5da415eb4df5e21960b564b40feb5e5c 100644 (file)
--- a/psycopg/psycopg/cursor_async.py
+++ b/psycopg/psycopg/cursor_async.py
@@ -212,6 +212,32 @@ class AsyncCursor(BaseCursor["AsyncConnection[Any]", Row]):
                 if not self.nextset():
                     break
 
+    async def set_result(self, index: int) -> Self:
+        """
+        Move to a specific result set.
+
+        :arg index: index of the result to go to
+        :type index: `!int`
+
+        More than one result will be available after executing calling
+        `executemany()` or `execute()` with more than one query.
+
+        `!index` is 0-based and supports negative values, counting from the end,
+        the same way you can index items in a list.
+
+        The function returns self, so that the result may be followed by a
+        fetch operation. See `results()` for details.
+        """
+        if not -len(self._results) <= index < len(self._results):
+            raise IndexError(
+                f"index {index} out of range: {len(self._results)} result(s) available"
+            )
+        if index < 0:
+            index = len(self._results) + index
+
+        self._select_current_result(index)
+        return self
+
     async def fetchone(self) -> Row | None:
         """
         Return the next record from the current result set.

diff --git a/tests/test_cursor_common.py b/tests/test_cursor_common.py
index 845e8e4e0a6a4c6a01f6565ddf7305036999c78f..c59650933f6cab646d6168dfd64f418fb5053ab9 100644 (file)
--- a/tests/test_cursor_common.py
+++ b/tests/test_cursor_common.py
@@ -220,23 +220,31 @@ def test_execute_many_results(conn):
     assert cur.rowcount == 3
     assert cur.nextset() is None
 
-    cur.set_result(0)
+    cur.close()
+    assert cur.nextset() is None
+
+
+def test_set_results(conn):
+    cur = conn.cursor()
+
+    with pytest.raises(IndexError):
+        cur.set_result(0)
+
+    cur.execute("select 'foo'; select generate_series(1,3)")
+    assert cur.set_result(0) is cur
     assert cur.fetchall() == [("foo",)]
     assert cur.rowcount == 1
 
-    cur.set_result(-1)
+    assert cur.set_result(-1) is cur
     assert cur.fetchall() == [(1,), (2,), (3,)]
     assert cur.rowcount == 3
 
-    with pytest.raises(ValueError):
+    with pytest.raises(IndexError):
         cur.set_result(2)
 
-    with pytest.raises(ValueError):
+    with pytest.raises(IndexError):
         cur.set_result(-3)
 
-    cur.close()
-    assert cur.nextset() is None
-
 
 def test_execute_sequence(conn):
     cur = conn.cursor()

diff --git a/tests/test_cursor_common_async.py b/tests/test_cursor_common_async.py
index 87765e623e2ca20bb1b648cfdcebd02c0dcba4ce..c565ab712474b3da07f931fa151896970def5fa8 100644 (file)
--- a/tests/test_cursor_common_async.py
+++ b/tests/test_cursor_common_async.py
@@ -218,22 +218,30 @@ async def test_execute_many_results(aconn):
     assert cur.rowcount == 3
     assert cur.nextset() is None
 
-    cur.set_result(0)
-    assert cur.fetchall() == [("foo",)]
+    await cur.close()
+    assert cur.nextset() is None
+
+
+async def test_set_results(aconn):
+    cur = aconn.cursor()
+
+    with pytest.raises(IndexError):
+        await cur.set_result(0)
+
+    await cur.execute("select 'foo'; select generate_series(1,3)")
+    assert await cur.set_result(0) is cur
+    assert (await cur.fetchall()) == [("foo",)]
     assert cur.rowcount == 1
 
-    cur.set_result(-1)
-    assert cur.fetchall() == [(1,), (2,), (3,)]
+    assert await cur.set_result(-1) is cur
+    assert (await cur.fetchall()) == [(1,), (2,), (3,)]
     assert cur.rowcount == 3
 
-    with pytest.raises(ValueError):
-        cur.set_result(2)
-
-    with pytest.raises(ValueError):
-        cur.set_result(-3)
+    with pytest.raises(IndexError):
+        await cur.set_result(2)
 
-    await cur.close()
-    assert cur.nextset() is None
+    with pytest.raises(IndexError):
+        await cur.set_result(-3)
 
 
 async def test_execute_sequence(aconn):