Catch PyStopAsyncIteration in callback_impl_loop_step macro

emmett-framework · Jul 9, 2024 · e0c3c4a · e0c3c4a
1 parent 24db55e
commit e0c3c4a
Showing 1 changed file with 10 additions and 9 deletions.
diff --git a/src/callbacks.rs b/src/callbacks.rs
@@ -357,7 +357,6 @@ macro_rules! callback_impl_loop_step {
                             ($pyself.into_py($py).getattr($py, pyo3::intern!($py, "_loop_wake"))?,),
                             Some(&kwctx),
                         )?;
-                        Ok(())
                     }
                     false => {
                         let event_loop = $pyself.context.event_loop($py);
@@ -366,18 +365,20 @@ macro_rules! callback_impl_loop_step {
                             ($pyself.into_py($py).getattr($py, pyo3::intern!($py, "_loop_step"))?,),
                             Some(&kwctx),
                         )?;
-                        Ok(())
                     }
                 }
+                Ok(())
             }
-            Err(err) => {
+            Err(err)
                 if (err.is_instance_of::<pyo3::exceptions::PyStopIteration>($py)
-                    || err.is_instance_of::<pyo3::exceptions::asyncio::CancelledError>($py))
-                {
-                    $pyself.done();
-                } else {
-                    $pyself.err(&err);
-                }
+                    || err.is_instance_of::<pyo3::exceptions::PyStopAsyncIteration>($py)
+                    || err.is_instance_of::<pyo3::exceptions::asyncio::CancelledError>($py)) =>
+            {
+                $pyself.done();
+                Ok(())
+            }
+            Err(err) => {
+                $pyself.err(&err);
                 Ok(())
             }
         }