Add set_object_tag

Author: corona10

Lib/test/test_sys.py

@@ -794,6 +794,20 @@ def test_get_object_tags(self):
         for k in keys:
             self.assertIn(k, tags)
 
+    @support.cpython_only
+    def test_set_object_tags(self):
+        keys = ("immortal", "interned")
+        s = "should never interned before" + str(random.randrange(0, 10**9))
+        origin_tags = sys.get_object_tags(s)
+        for k in keys:
+            sys.set_object_tag(s, k)
+        sys.set_object_tag(s, "unknown")
+        after_tags = sys.get_object_tags(s)
+        self.assertEqual(len(origin_tags), len(after_tags))
+        for k in keys:
+            self.assertIn(k, after_tags)
+            self.assertTrue(after_tags[k])
+
     @support.cpython_only
     @requires_subinterpreters
     def test_subinterp_intern_dynamically_allocated(self):

Python/clinic/sysmodule.c.h

@@ -1083,7 +1083,7 @@ sys_get_object_tags(PyObject *module, PyObject *op)
         }
     }
 
-    if (PyUnicode_CHECK_INTERNED(op)) {
+    if (PyUnicode_Check(op) && PyUnicode_CHECK_INTERNED(op)) {
         if (PyDict_SetItemString(dict, "interned", Py_True) < 0) {
             Py_DECREF(dict);
             return NULL;
@@ -1111,6 +1111,32 @@ sys_get_object_tags(PyObject *module, PyObject *op)
     return dict;
 }
 
+/*[clinic input]
+sys.set_object_tag -> object
+
+  object: object
+  tag: str
+  *
+  options: object = None
+
+Set the tags of the given object.
+[clinic start generated code]*/
+
+static PyObject *
+sys_set_object_tag_impl(PyObject *module, PyObject *object, const char *tag,
+                        PyObject *options)
+/*[clinic end generated code: output=b0fb5e9931feb4aa input=b64c9bd958c75f11]*/
+{
+    assert(object != NULL);
+    if (strcmp(tag, "immortal") == 0) {
+        _Py_SetImmortal(object);
+    }
+    else if (strcmp(tag, "interned") == 0) {
+        _PyUnicode_InternMortal(_PyInterpreterState_GET(), &object);
+    }
+    Py_RETURN_NONE;
+}
+
 /*
  * Cached interned string objects used for calling the profile and
  * trace functions.
@@ -2856,6 +2882,7 @@ static PyMethodDef sys_methods[] = {
     SYS_INTERN_METHODDEF
     SYS__IS_INTERNED_METHODDEF
     SYS_GET_OBJECT_TAGS_METHODDEF
+    SYS_SET_OBJECT_TAG_METHODDEF
     SYS_IS_FINALIZING_METHODDEF
     SYS_MDEBUG_METHODDEF
     SYS_SETSWITCHINTERVAL_METHODDEF