Use the new PyAwaitable_AddExpr function inside of PyAwaitable_AsyncWith.

This should simplify its implementation and helps eliminate any duplicate code.

Signed-off-by: AraHaan <[email protected]>

Author: AraHaan

src/_pyawaitable/with.c

@@ -40,44 +40,43 @@ async_with_inner(PyObject *aw, PyObject *res)
         if (tb == NULL)
             tb = Py_NewRef(Py_None);
 
-        PyObject *coro = PyObject_Vectorcall(
-            exit,
-            (PyObject *[]) { tp, val, tb },
-            3,
-            NULL
-        );
-        Py_DECREF(tp);
-        Py_DECREF(val);
-        Py_DECREF(tb);
-        if (coro == NULL) {
+        if (
+            PyAwaitable_AddExpr(
+                aw,
+                PyObject_Vectorcall(
+                    exit,
+                    (PyObject *[]) { tp, val, tb },
+                    3,
+                    NULL),
+                NULL,
+                NULL) < 0
+        ) {
+            Py_DECREF(tp);
+            Py_DECREF(val);
+            Py_DECREF(tb);
             return -1;
         }
 
-        if (PyAwaitable_AddAwait(aw, coro, NULL, NULL) < 0) {
-            Py_DECREF(coro);
-            return -1;
-        }
-
-        Py_DECREF(coro);
+        Py_DECREF(tp);
+        Py_DECREF(val);
+        Py_DECREF(tb);
         return 0;
     }
     else {
         // OK
-        PyObject *coro = PyObject_Vectorcall(
-            exit,
-            (PyObject *[]) { Py_None, Py_None, Py_None },
-            3,
-            NULL
-        );
-        if (coro == NULL) {
+        if (
+            PyAwaitable_AddExpr(
+                aw,
+                PyObject_Vectorcall(
+                    exit,
+                    (PyObject *[]) { Py_None, Py_None, Py_None },
+                    3,
+                    NULL),
+                NULL,
+                NULL) < 0
+        ) {
             return -1;
         }
-
-        if (PyAwaitable_AddAwait(aw, coro, NULL, NULL) < 0) {
-            Py_DECREF(coro);
-            return -1;
-        }
-        Py_DECREF(coro);
         return 0;
     }
 }
@@ -134,35 +133,25 @@ PyAwaitable_AsyncWith(
 
     Py_DECREF(exit);
 
-    PyObject *coro = PyObject_CallNoArgs(with);
-    Py_DECREF(with);
-
-    if (coro == NULL) {
-        Py_DECREF(inner_aw);
-        return -1;
-    }
-
     // Note: Errors in __aenter__ are not sent to __aexit__
     if (
-        PyAwaitable_AddAwait(
+        PyAwaitable_AddExpr(
             inner_aw,
-            coro,
+            PyObject_CallNoArgs(with),
             async_with_inner,
             NULL
         ) < 0
     ) {
         Py_DECREF(inner_aw);
-        Py_DECREF(coro);
+        Py_DECREF(with);
         return -1;
     }
 
-    Py_DECREF(coro);
+    Py_DECREF(with);
 
-    if (PyAwaitable_AddAwait(aw, inner_aw, NULL, err) < 0) {
-        Py_DECREF(inner_aw);
+    if (PyAwaitable_AddExpr(aw, inner_aw, NULL, err) < 0) {
         return -1;
     }
 
-    Py_DECREF(inner_aw);
     return 0;
 }