From 942916378aa6a0946b1385c2c7ca6935620d710a Mon Sep 17 00:00:00 2001
From: Victor Stinner <vstinner@python.org>
Date: Thu, 10 Oct 2024 01:37:14 +0200
Subject: [PATCH] gh-125196: Use PyUnicodeWriter for repr(contextvars.Token)
 (#125220)

Replace the private _PyUnicodeWriter with the public PyUnicodeWriter.
---
 Lib/test/test_context.py |  8 ++++++++
 Python/context.c         | 39 +++++++++++----------------------------
 2 files changed, 19 insertions(+), 28 deletions(-)

diff --git a/Lib/test/test_context.py b/Lib/test/test_context.py
index 255be306156c0b..b06b9df9f5b0b8 100644
--- a/Lib/test/test_context.py
+++ b/Lib/test/test_context.py
@@ -60,6 +60,14 @@ def test_context_var_repr_1(self):
         c.reset(t)
         self.assertIn(' used ', repr(t))
 
+    @isolated_context
+    def test_token_repr_1(self):
+        c = contextvars.ContextVar('a')
+        tok = c.set(1)
+        self.assertRegex(repr(tok),
+                         r"^<Token var=<ContextVar name='a' "
+                         r"at 0x[0-9a-fA-F]+> at 0x[0-9a-fA-F]+>$")
+
     def test_context_subclassing_1(self):
         with self.assertRaisesRegex(TypeError, 'not an acceptable base type'):
             class MyContextVar(contextvars.ContextVar):
diff --git a/Python/context.c b/Python/context.c
index ddb03555f9e402..36e2677c398f59 100644
--- a/Python/context.c
+++ b/Python/context.c
@@ -1154,48 +1154,31 @@ token_tp_dealloc(PyContextToken *self)
 static PyObject *
 token_tp_repr(PyContextToken *self)
 {
-    _PyUnicodeWriter writer;
-
-    _PyUnicodeWriter_Init(&writer);
-
-    if (_PyUnicodeWriter_WriteASCIIString(&writer, "<Token", 6) < 0) {
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create(0);
+    if (writer == NULL) {
+        return NULL;
+    }
+    if (PyUnicodeWriter_WriteUTF8(writer, "<Token", 6) < 0) {
         goto error;
     }
-
     if (self->tok_used) {
-        if (_PyUnicodeWriter_WriteASCIIString(&writer, " used", 5) < 0) {
+        if (PyUnicodeWriter_WriteUTF8(writer, " used", 5) < 0) {
             goto error;
         }
     }
-
-    if (_PyUnicodeWriter_WriteASCIIString(&writer, " var=", 5) < 0) {
+    if (PyUnicodeWriter_WriteUTF8(writer, " var=", 5) < 0) {
         goto error;
     }
-
-    PyObject *var = PyObject_Repr((PyObject *)self->tok_var);
-    if (var == NULL) {
+    if (PyUnicodeWriter_WriteRepr(writer, (PyObject *)self->tok_var) < 0) {
         goto error;
     }
-    if (_PyUnicodeWriter_WriteStr(&writer, var) < 0) {
-        Py_DECREF(var);
-        goto error;
-    }
-    Py_DECREF(var);
-
-    PyObject *addr = PyUnicode_FromFormat(" at %p>", self);
-    if (addr == NULL) {
-        goto error;
-    }
-    if (_PyUnicodeWriter_WriteStr(&writer, addr) < 0) {
-        Py_DECREF(addr);
+    if (PyUnicodeWriter_Format(writer, " at %p>", self) < 0) {
         goto error;
     }
-    Py_DECREF(addr);
-
-    return _PyUnicodeWriter_Finish(&writer);
+    return PyUnicodeWriter_Finish(writer);
 
 error:
-    _PyUnicodeWriter_Dealloc(&writer);
+    PyUnicodeWriter_Discard(writer);
     return NULL;
 }