[pypy-commit] pypy default: Expose a few members of the PyCodeObject structure

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: 
Changeset: r47176:9962794cb329
Date: 2011-09-08 01:04 +0200
http://bitbucket.org/pypy/pypy/changeset/9962794cb329/

Log:	Expose a few members of the PyCodeObject structure

diff --git a/pypy/interpreter/pycode.py b/pypy/interpreter/pycode.py
--- a/pypy/interpreter/pycode.py
+++ b/pypy/interpreter/pycode.py
@@ -10,7 +10,7 @@
 from pypy.interpreter.argument import Signature
 from pypy.interpreter.error import OperationError
 from pypy.interpreter.gateway import NoneNotWrapped, unwrap_spec
-from pypy.interpreter.astcompiler.consts import (CO_OPTIMIZED,
+from pypy.interpreter.astcompiler.consts import (
     CO_OPTIMIZED, CO_NEWLOCALS, CO_VARARGS, CO_VARKEYWORDS, CO_NESTED,
     CO_GENERATOR, CO_CONTAINSGLOBALS)
 from pypy.rlib.rarithmetic import intmask
diff --git a/pypy/module/cpyext/funcobject.py b/pypy/module/cpyext/funcobject.py
--- a/pypy/module/cpyext/funcobject.py
+++ b/pypy/module/cpyext/funcobject.py
@@ -4,9 +4,21 @@
     cpython_api, bootstrap_function, cpython_struct, build_type_checkers)
 from pypy.module.cpyext.pyobject import (
     PyObject, make_ref, from_ref, Py_DecRef, make_typedescr, borrow_from)
+from pypy.rlib.unroll import unrolling_iterable
 from pypy.interpreter.error import OperationError
 from pypy.interpreter.function import Function, Method
 from pypy.interpreter.pycode import PyCode
+from pypy.interpreter import pycode
+
+CODE_FLAGS = dict(
+    CO_OPTIMIZED   = 0x0001,
+    CO_NEWLOCALS   = 0x0002,
+    CO_VARARGS     = 0x0004,
+    CO_VARKEYWORDS = 0x0008,
+    CO_NESTED      = 0x0010,
+    CO_GENERATOR   = 0x0020,
+)
+ALL_CODE_FLAGS = unrolling_iterable(CODE_FLAGS.items())
 
 PyFunctionObjectStruct = lltype.ForwardReference()
 PyFunctionObject = lltype.Ptr(PyFunctionObjectStruct)
@@ -16,7 +28,12 @@
 
 PyCodeObjectStruct = lltype.ForwardReference()
 PyCodeObject = lltype.Ptr(PyCodeObjectStruct)
-cpython_struct("PyCodeObject", PyObjectFields, PyCodeObjectStruct)
+PyCodeObjectFields = PyObjectFields + \
+    (("co_name", PyObject),
+     ("co_flags", rffi.INT),
+     ("co_argcount", rffi.INT),
+    )
+cpython_struct("PyCodeObject", PyCodeObjectFields, PyCodeObjectStruct)
 
 @bootstrap_function
 def init_functionobject(space):
@@ -24,6 +41,10 @@
                    basestruct=PyFunctionObject.TO,
                    attach=function_attach,
                    dealloc=function_dealloc)
+    make_typedescr(PyCode.typedef,
+                   basestruct=PyCodeObject.TO,
+                   attach=code_attach,
+                   dealloc=code_dealloc)
 
 PyFunction_Check, PyFunction_CheckExact = build_type_checkers("Function", Function)
 PyMethod_Check, PyMethod_CheckExact = build_type_checkers("Method", Method)
@@ -40,6 +61,24 @@
     from pypy.module.cpyext.object import PyObject_dealloc
     PyObject_dealloc(space, py_obj)
 
+def code_attach(space, py_obj, w_obj):
+    py_code = rffi.cast(PyCodeObject, py_obj)
+    assert isinstance(w_obj, PyCode)
+    py_code.c_co_name = make_ref(space, space.wrap(w_obj.co_name))
+    co_flags = 0
+    for name, value in ALL_CODE_FLAGS:
+        if w_obj.co_flags & getattr(pycode, name):
+            co_flags |= value
+    rffi.setintfield(py_code, 'c_co_flags', co_flags)
+    rffi.setintfield(py_code, 'c_co_argcount', w_obj.co_argcount)
+
+ at cpython_api([PyObject], lltype.Void, external=False)
+def code_dealloc(space, py_obj):
+    py_code = rffi.cast(PyCodeObject, py_obj)
+    Py_DecRef(space, py_code.c_co_name)
+    from pypy.module.cpyext.object import PyObject_dealloc
+    PyObject_dealloc(space, py_obj)
+
 @cpython_api([PyObject], PyObject)
 def PyFunction_GetCode(space, w_func):
     """Return the code object associated with the function object op."""
diff --git a/pypy/module/cpyext/include/code.h b/pypy/module/cpyext/include/code.h
--- a/pypy/module/cpyext/include/code.h
+++ b/pypy/module/cpyext/include/code.h
@@ -4,7 +4,21 @@
 extern "C" {
 #endif
 
-typedef PyObject PyCodeObject;
+typedef struct {
+    PyObject_HEAD
+    PyObject *co_name;
+    int co_argcount;
+    int co_flags;
+} PyCodeObject;
+
+/* Masks for co_flags above */
+/* These values are also in funcobject.py */
+#define CO_OPTIMIZED	0x0001
+#define CO_NEWLOCALS	0x0002
+#define CO_VARARGS	0x0004
+#define CO_VARKEYWORDS	0x0008
+#define CO_NESTED       0x0010
+#define CO_GENERATOR    0x0020
 
 #ifdef __cplusplus
 }
diff --git a/pypy/module/cpyext/test/test_funcobject.py b/pypy/module/cpyext/test/test_funcobject.py
--- a/pypy/module/cpyext/test/test_funcobject.py
+++ b/pypy/module/cpyext/test/test_funcobject.py
@@ -2,8 +2,12 @@
 from pypy.module.cpyext.test.test_cpyext import AppTestCpythonExtensionBase
 from pypy.module.cpyext.test.test_api import BaseApiTest
 from pypy.module.cpyext.pyobject import PyObject, make_ref, from_ref
-from pypy.module.cpyext.funcobject import PyFunctionObject
+from pypy.module.cpyext.funcobject import (
+    PyFunctionObject, PyCodeObject, CODE_FLAGS)
 from pypy.interpreter.function import Function, Method
+from pypy.interpreter.pycode import PyCode
+
+globals().update(CODE_FLAGS)
 
 class TestFunctionObject(BaseApiTest):
     def test_function(self, space, api):
@@ -38,12 +42,36 @@
 
     def test_getcode(self, space, api):
         w_function = space.appexec([], """():
-            def func(x): return x
+            def func(x, y, z): return x
             return func
         """)
         w_code = api.PyFunction_GetCode(w_function)
         assert w_code.co_name == "func"
 
+        ref = make_ref(space, w_code)
+        assert (from_ref(space, rffi.cast(PyObject, ref.c_ob_type)) is
+                space.gettypeobject(PyCode.typedef))
+        assert "func" == space.unwrap(
+           from_ref(space, rffi.cast(PyCodeObject, ref).c_co_name))
+        assert 3 == rffi.cast(PyCodeObject, ref).c_co_argcount
+        api.Py_DecRef(ref)
+
+    def test_co_flags(self, space, api):
+        def get_flags(signature, body="pass"):
+            w_code = space.appexec([], """():
+                def func(%s): %s
+                return func.__code__
+            """ % (signature, body))
+            ref = make_ref(space, w_code)
+            co_flags = rffi.cast(PyCodeObject, ref).c_co_flags
+            api.Py_DecRef(ref)
+            return co_flags
+        assert get_flags("x") == CO_NESTED | CO_OPTIMIZED | CO_NEWLOCALS
+        assert get_flags("x", "exec x") == CO_NESTED | CO_NEWLOCALS
+        assert get_flags("x, *args") & CO_VARARGS
+        assert get_flags("x, **kw") & CO_VARKEYWORDS
+        assert get_flags("x", "yield x") & CO_GENERATOR
+
     def test_newcode(self, space, api):
         filename = rffi.str2charp('filename')
         funcname = rffi.str2charp('funcname')