[Python-checkins] r50738 - in python/branches/bcannon-sandboxing: Lib/test/test_sys.py Python/sysmodule.c

[brett.cannon]

Author: brett.cannon
Date: Thu Jul 20 23:35:05 2006
New Revision: 50738

Modified:
   python/branches/bcannon-sandboxing/Lib/test/test_sys.py
   python/branches/bcannon-sandboxing/Python/sysmodule.c
Log:
Add support to set/get memory cap and get memory usage in the sys module.


Modified: python/branches/bcannon-sandboxing/Lib/test/test_sys.py
==============================================================================
--- python/branches/bcannon-sandboxing/Lib/test/test_sys.py	(original)
+++ python/branches/bcannon-sandboxing/Lib/test/test_sys.py	Thu Jul 20 23:35:05 2006
@@ -266,6 +266,25 @@
         # the test runs under regrtest.
         self.assert_(sys.__stdout__.encoding == sys.__stderr__.encoding)
 
+    def test_memorycap(self):
+        # Test both getmemorycap() and setmemorycap().
+        if not hasattr(sys, 'getmemorycap'):
+            return
+        original_cap = sys.getmemorycap()
+        self.failUnless(isinstance(original_cap, (int, long)))
+        new_cap = int(original_cap + 10000)
+        assert isinstance(new_cap, int)
+        sys.setmemorycap(new_cap)
+        try:  # Make sure don't mess up interpreter.
+            self.failUnlessEqual(new_cap, sys.getmemorycap())
+            sys.setmemorycap(long(new_cap))
+            self.failUnlessEqual(new_cap, sys.getmemorycap())
+        finally:
+            try:  # setmemorycap() could be broken.
+                sys.setmemorycap(original_cap)
+            except Exception:
+                pass
+
 def test_main():
     test.test_support.run_unittest(SysModuleTest)
 

Modified: python/branches/bcannon-sandboxing/Python/sysmodule.c
==============================================================================
--- python/branches/bcannon-sandboxing/Python/sysmodule.c	(original)
+++ python/branches/bcannon-sandboxing/Python/sysmodule.c	Thu Jul 20 23:35:05 2006
@@ -700,6 +700,64 @@
 10. Number of stack pops performed by call_function()"
 );
 
+#ifdef Py_MEMORY_CAP
+static PyObject *
+sys_setmemorycap(PyObject *self, PyObject *arg)
+{
+    PyInterpreterState *interp = PyInterpreterState_SafeGet();
+    PY_LONG_LONG new_memory_cap;
+    PyObject *arg_as_long = PyNumber_Long(arg);
+
+    if (!arg_as_long)
+	return NULL;
+
+    new_memory_cap = PyLong_AsLongLong(arg_as_long);
+    Py_DECREF(arg_as_long); /* DEAD: arg_as_long */
+
+    if (!interp)
+	Py_FatalError("interpreter not available");
+
+    if (!PyInterpreterState_SetMemoryCap(interp, new_memory_cap))
+	return NULL;
+
+    Py_RETURN_NONE;
+}
+
+PyDoc_STRVAR(setmemorycap_doc,
+"XXX"
+);
+
+static PyObject *
+sys_getmemorycap(PyObject *self, PyObject *ignore)
+{
+    PyInterpreterState *interp = PyInterpreterState_SafeGet();
+
+    if (!interp)
+	Py_FatalError("interpreter not available");
+
+    return PyLong_FromLongLong(interp->mem_cap);
+}
+
+PyDoc_STRVAR(getmemorycap_doc,
+"XXX"
+);
+
+static PyObject *
+sys_getmemoryused(PyObject *self, PyObject *ignore)
+{
+    PyInterpreterState *interp = PyInterpreterState_SafeGet();
+
+    if (!interp)
+	Py_FatalError("interpreter not available");
+
+    return PyLong_FromLongLong(interp->mem_usage);
+}
+
+PyDoc_STRVAR(getmemoryused_doc,
+"XXX"
+);
+#endif /* Py_MEMORY_CAP */
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -782,6 +840,11 @@
 #endif
 	{"settrace",	sys_settrace, METH_O, settrace_doc},
 	{"call_tracing", sys_call_tracing, METH_VARARGS, call_tracing_doc},
+#ifdef Py_MEMORY_CAP
+	{"setmemorycap", sys_setmemorycap, METH_O, setmemorycap_doc},
+	{"getmemorycap", sys_getmemorycap, METH_NOARGS, getmemorycap_doc},
+	{"getmemoryused", sys_getmemoryused, METH_NOARGS, getmemoryused_doc},
+#endif 
 	{NULL,		NULL}		/* sentinel */
 };