diff --git a/Lib/test/test_types.py b/Lib/test/test_types.py index fc26e71ffcb67b..c1c5bc0d700e72 100644 --- a/Lib/test/test_types.py +++ b/Lib/test/test_types.py @@ -21,6 +21,8 @@ import unittest.mock import weakref import typing +from types import SimpleNamespace + c_types = import_fresh_module('types', fresh=['_types']) py_types = import_fresh_module('types', blocked=['_types']) @@ -2128,6 +2130,23 @@ class FakeSimpleNamespace(str): types.SimpleNamespace() > FakeSimpleNamespace() with self.assertRaises(TypeError): types.SimpleNamespace() >= FakeSimpleNamespace() + + def test_update_method(self): + ns = SimpleNamespace(a=1) + self.assertEqual(ns.a, 1) + + ns.update(b=2, c=3) + self.assertEqual(ns.b, 2) + self.assertEqual(ns.c, 3) + + # Overwriting existing key + ns.update(a=42) + self.assertEqual(ns.a, 42) + + # No update with no kwargs + ns.update() + self.assertEqual(vars(ns), {'a': 42, 'b': 2, 'c': 3}) + class CoroutineTests(unittest.TestCase): diff --git a/Objects/namespaceobject.c b/Objects/namespaceobject.c index 0fc2bcea4cb06e..52b52094035b51 100644 --- a/Objects/namespaceobject.c +++ b/Objects/namespaceobject.c @@ -245,6 +245,37 @@ namespace_replace(PyObject *self, PyObject *args, PyObject *kwargs) return result; } +static PyObject *namespace_update(PyObject *self, PyObject *args, PyObject *kwds); +static PyObject * +namespace_update(PyObject *self, PyObject *args, PyObject *kwargs) +{ + PyObject *update_obj = NULL; + + if (!PyArg_UnpackTuple(args, "update", 0, 1, &update_obj)) { + return NULL; + } + + _PyNamespaceObject *ns = (_PyNamespaceObject *)self; + PyObject *dict = ns->ns_dict; + Py_INCREF(dict); + + if (update_obj != NULL) { + if (PyDict_Update(dict, update_obj) < 0) { + Py_DECREF(dict); + return NULL; + } + } + + if (kwargs && PyDict_Update(dict, kwargs) < 0) { + Py_DECREF(dict); + return NULL; + } + + Py_DECREF(dict); + Py_RETURN_NONE; +} + + static PyMethodDef namespace_methods[] = { {"__reduce__", namespace_reduce, METH_NOARGS, @@ -252,7 +283,15 @@ static PyMethodDef namespace_methods[] = { {"__replace__", _PyCFunction_CAST(namespace_replace), METH_VARARGS|METH_KEYWORDS, PyDoc_STR("__replace__($self, /, **changes)\n--\n\n" "Return a copy of the namespace object with new values for the specified attributes.")}, - {NULL, NULL} // sentinel + + {"update", (PyCFunction)(PyCFunctionWithKeywords)namespace_update, METH_VARARGS | METH_KEYWORDS, + PyDoc_STR("update(self, mapping=(), /, **kwargs)\n--\n\n" + "Update the namespace with the given keyword arguments.")}, + + {NULL, NULL} /* sentinel */ + + + }; @@ -321,3 +360,8 @@ _PyNamespace_New(PyObject *kwds) return (PyObject *)ns; } + + + + +