From: Denis Laxalde <denis@laxalde.org>
Date: Thu, 11 Nov 2021 09:59:59 +0000 (+0100)
Subject: Handle RW-ready in wait(,_conn)_async()
X-Git-Tag: 3.0.5~10^2~2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=84a438e8d483a6258cacf4a351dc07ac6f4e8fee;p=thirdparty%2Fpsycopg.git

Handle RW-ready in wait(,_conn)_async()
---

diff --git a/psycopg/psycopg/waiting.py b/psycopg/psycopg/waiting.py
index b75dbe5b5..7563d4f44 100644
--- a/psycopg/psycopg/waiting.py
+++ b/psycopg/psycopg/waiting.py
@@ -119,29 +119,27 @@ async def wait_async(gen: PQGen[RV], fileno: int) -> RV:
 
     def wakeup(state: Ready) -> None:
         nonlocal ready
-        ready = state
+        ready |= state  # type: ignore[assignment]
         ev.set()
 
     try:
         s = next(gen)
         while 1:
+            reader = s & Wait.R
+            writer = s & Wait.W
+            if not reader and not writer:
+                raise e.InternalError(f"bad poll status: {s}")
             ev.clear()
-            if s == Wait.R:
-                loop.add_reader(fileno, wakeup, Ready.R)
-                await ev.wait()
-                loop.remove_reader(fileno)
-            elif s == Wait.W:
-                loop.add_writer(fileno, wakeup, Ready.W)
-                await ev.wait()
-                loop.remove_writer(fileno)
-            elif s == Wait.RW:
+            ready = 0  # type: ignore[assignment]
+            if reader:
                 loop.add_reader(fileno, wakeup, Ready.R)
+            if writer:
                 loop.add_writer(fileno, wakeup, Ready.W)
-                await ev.wait()
+            await ev.wait()
+            if reader:
                 loop.remove_reader(fileno)
+            if writer:
                 loop.remove_writer(fileno)
-            else:
-                raise e.InternalError("bad poll status: %s")
             s = gen.send(ready)
 
     except StopIteration as ex:
@@ -180,23 +178,21 @@ async def wait_conn_async(
     try:
         fileno, s = next(gen)
         while 1:
+            reader = s & Wait.R
+            writer = s & Wait.W
+            if not reader and not writer:
+                raise e.InternalError(f"bad poll status: {s}")
             ev.clear()
-            if s == Wait.R:
-                loop.add_reader(fileno, wakeup, Ready.R)
-                await wait_for(ev.wait(), timeout)
-                loop.remove_reader(fileno)
-            elif s == Wait.W:
-                loop.add_writer(fileno, wakeup, Ready.W)
-                await wait_for(ev.wait(), timeout)
-                loop.remove_writer(fileno)
-            elif s == Wait.RW:
+            ready = 0  # type: ignore[assignment]
+            if reader:
                 loop.add_reader(fileno, wakeup, Ready.R)
+            if writer:
                 loop.add_writer(fileno, wakeup, Ready.W)
-                await wait_for(ev.wait(), timeout)
+            await wait_for(ev.wait(), timeout)
+            if reader:
                 loop.remove_reader(fileno)
+            if writer:
                 loop.remove_writer(fileno)
-            else:
-                raise e.InternalError("bad poll status: %s")
             fileno, s = gen.send(ready)
 
     except TimeoutError: