1 files changed, 21 insertions, 10 deletions

diff --git a/Python/context.c b/Python/context.c
index 8bc487a..95aa822 100644
--- a/Python/context.c
+++ b/Python/context.c
@@ -102,10 +102,8 @@ PyContext_CopyCurrent(void)
 static const char *
 context_event_name(PyContextEvent event) {
     switch (event) {
-        case Py_CONTEXT_EVENT_ENTER:
-            return "Py_CONTEXT_EVENT_ENTER";
-        case Py_CONTEXT_EVENT_EXIT:
-            return "Py_CONTEXT_EVENT_EXIT";
+        case Py_CONTEXT_SWITCHED:
+            return "Py_CONTEXT_SWITCHED";
         default:
             return "?";
     }
@@ -115,6 +113,13 @@ context_event_name(PyContextEvent event) {
 static void
 notify_context_watchers(PyThreadState *ts, PyContextEvent event, PyObject *ctx)
 {
+    if (ctx == NULL) {
+        // This will happen after exiting the last context in the stack, which
+        // can occur if context_get was never called before entering a context
+        // (e.g., called `contextvars.Context().run()` on a fresh thread, as
+        // PyContext_Enter doesn't call context_get).
+        ctx = Py_None;
+    }
     assert(Py_REFCNT(ctx) > 0);
     PyInterpreterState *interp = ts->interp;
     assert(interp->_initialized);
@@ -175,6 +180,16 @@ PyContext_ClearWatcher(int watcher_id)
 }
 
 
+static inline void
+context_switched(PyThreadState *ts)
+{
+    ts->context_ver++;
+    // ts->context is used instead of context_get() because context_get() might
+    // throw if ts->context is NULL.
+    notify_context_watchers(ts, Py_CONTEXT_SWITCHED, ts->context);
+}
+
+
 static int
 _PyContext_Enter(PyThreadState *ts, PyObject *octx)
 {
@@ -191,9 +206,7 @@ _PyContext_Enter(PyThreadState *ts, PyObject *octx)
     ctx->ctx_entered = 1;
 
     ts->context = Py_NewRef(ctx);
-    ts->context_ver++;
-
-    notify_context_watchers(ts, Py_CONTEXT_EVENT_ENTER, octx);
+    context_switched(ts);
     return 0;
 }
 
@@ -227,13 +240,11 @@ _PyContext_Exit(PyThreadState *ts, PyObject *octx)
         return -1;
     }
 
-    notify_context_watchers(ts, Py_CONTEXT_EVENT_EXIT, octx);
     Py_SETREF(ts->context, (PyObject *)ctx->ctx_prev);
-    ts->context_ver++;
 
     ctx->ctx_prev = NULL;
     ctx->ctx_entered = 0;
-
+    context_switched(ts);
     return 0;
 }