Using ctypes to access the type object, set object dictionaries

pyctypes.py

'''
Created on 27.02.2021

@author: anselm
'''

# Some code copied from
#
# forbiddenfruit - Patch built-in python objects
#
# Copyright (c) 2013-2020  Lincoln de Sousa <[email protected]>
#
# MIT
# ---
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import sys
import ctypes

Py_ssize_t = ctypes.c_ssize_t


# # dictionary holding references to the allocated function resolution
# # arrays to type objects
# tp_as_dict = {}
# # container to cfunc callbacks
# tp_func_dict = {}


class PyObject(ctypes.Structure):
    pass

class PyVarObject(ctypes.Structure):
    pass

class PyFile(ctypes.Structure):
    pass

PyObject_p = ctypes.py_object
Inquiry_p = ctypes.CFUNCTYPE(ctypes.c_int, PyObject_p)
# return type is void* to allow ctypes to convert python integers to
# plain PyObject*
UnaryFunc_p = ctypes.CFUNCTYPE(ctypes.py_object, PyObject_p)
BinaryFunc_p = ctypes.CFUNCTYPE(ctypes.py_object, PyObject_p, PyObject_p)
TernaryFunc_p = ctypes.CFUNCTYPE(ctypes.py_object, PyObject_p, PyObject_p, PyObject_p)
LenFunc_p = ctypes.CFUNCTYPE(Py_ssize_t, PyObject_p)
SSizeArgFunc_p = ctypes.CFUNCTYPE(ctypes.py_object, PyObject_p, Py_ssize_t)
SSizeObjArgProc_p = ctypes.CFUNCTYPE(ctypes.c_int, PyObject_p, Py_ssize_t, PyObject_p)
ObjObjProc_p = ctypes.CFUNCTYPE(ctypes.c_int, PyObject_p, PyObject_p)
Destructor_p = ctypes.CFUNCTYPE(None, PyObject_p)

FILE_p = ctypes.POINTER(PyFile)


class PyNumberMethods(ctypes.Structure):
    _fields_ = [
    ('nb_add', BinaryFunc_p),
    ('nb_subtract', BinaryFunc_p),
    ('nb_multiply', BinaryFunc_p),
    ('nb_remainder', BinaryFunc_p),
    ('nb_divmod', BinaryFunc_p),
    ('nb_power', BinaryFunc_p),
    ('nb_negative', UnaryFunc_p),
    ('nb_positive', UnaryFunc_p),
    ('nb_absolute', UnaryFunc_p),
    ('nb_bool', Inquiry_p),
    ('nb_invert', UnaryFunc_p),
    ('nb_lshift', BinaryFunc_p),
    ('nb_rshift', BinaryFunc_p),
    ('nb_and', BinaryFunc_p),
    ('nb_xor', BinaryFunc_p),
    ('nb_or', BinaryFunc_p),
    ('nb_int', UnaryFunc_p),
    ('nb_reserved', ctypes.c_void_p),
    ('nb_float', UnaryFunc_p),

    ('nb_inplace_add', BinaryFunc_p),
    ('nb_inplace_subtract', BinaryFunc_p),
    ('nb_inplace_multiply', BinaryFunc_p),
    ('nb_inplace_remainder', BinaryFunc_p),
    ('nb_inplace_power', TernaryFunc_p),
    ('nb_inplace_lshift', BinaryFunc_p),
    ('nb_inplace_rshift', BinaryFunc_p),
    ('nb_inplace_and', BinaryFunc_p),
    ('nb_inplace_xor', BinaryFunc_p),
    ('nb_inplace_or', BinaryFunc_p),

    ('nb_floor_divide', BinaryFunc_p),
    ('nb_true_divide', BinaryFunc_p),
    ('nb_inplace_floor_divide', BinaryFunc_p),
    ('nb_inplace_true_divide', BinaryFunc_p),

    ('nb_index', BinaryFunc_p),

    ('nb_matrix_multiply', BinaryFunc_p),
    ('nb_inplace_matrix_multiply', BinaryFunc_p),
    ]

class PySequenceMethods(ctypes.Structure):
    _fields_ = [
        ('sq_length', LenFunc_p),
        ('sq_concat', BinaryFunc_p),
        ('sq_repeat', SSizeArgFunc_p),
        ('sq_item', SSizeArgFunc_p),
        ('was_sq_slice', ctypes.c_void_p),
        ('sq_ass_item', SSizeObjArgProc_p),
        ('was_sq_ass_slice', ctypes.c_void_p),
        ('sq_contains', ObjObjProc_p),
        ('sq_inplace_concat', BinaryFunc_p),
        ('sq_inplace_repeat', SSizeArgFunc_p),
    ]

class PyMappingMethods(ctypes.Structure):
    pass

class PyTypeObject(ctypes.Structure):
    pass

class PyAsyncMethods(ctypes.Structure):
    pass


class _PyObject_API(ctypes.Structure):
    _fields_ = [
        ('ob_refcnt', Py_ssize_t),
        ('ob_type', ctypes.POINTER(PyTypeObject)),
        ]

_PyObject_HEAD_size = object().__sizeof__()
# _PyVarObject_HEAD_size = ().__sizeof__()
# _PyGC_HEAD_size = sys.getsizeof(()) - _PyVarObject_HEAD_size

assert _PyObject_HEAD_size >= ctypes.sizeof(_PyObject_API) 

PyObject._fields_ = ([ 
    ('_HEAD_EXTRA', ctypes.c_void_p * ( (_PyObject_HEAD_size - ctypes.sizeof(_PyObject_API)) // ctypes.sizeof(ctypes.c_void_p)) ) 
] if _PyObject_HEAD_size > ctypes.sizeof(_PyObject_API) else []) + [
    ('ob_refcnt', Py_ssize_t),
    ('ob_type', ctypes.POINTER(PyTypeObject)),
]

assert ctypes.sizeof(PyObject) == _PyObject_HEAD_size

PyVarObject._fields_ = [ 
    ('ob_base', PyObject),
    ('ob_size', Py_ssize_t),
]

assert ctypes.sizeof(PyVarObject) == ().__sizeof__()

PyTypeObject._fields_ = [
    # varhead
    ('ob_base', PyVarObject),
    # declaration
    ('tp_name', ctypes.c_char_p),
    ('tp_basicsize', Py_ssize_t),
    ('tp_itemsize', Py_ssize_t),
    ('tp_dealloc', ctypes.CFUNCTYPE(None, PyObject_p)),
    (('tp_vectorcall_offset', Py_ssize_t) if sys.hexversion >= 0x030800b1 else 
    ('tp_print', ctypes.CFUNCTYPE(ctypes.c_int, PyObject_p, FILE_p, ctypes.c_int))),
    ('tp_getattr', ctypes.CFUNCTYPE(PyObject_p, PyObject_p, ctypes.c_char_p)),
    ('tp_setattr', ctypes.CFUNCTYPE(ctypes.c_int, PyObject_p, ctypes.c_char_p, PyObject_p)),
    ('tp_as_async', ctypes.CFUNCTYPE(PyAsyncMethods)),
    ('tp_repr', ctypes.CFUNCTYPE(PyObject_p, PyObject_p)),
    ('tp_as_number', ctypes.POINTER(PyNumberMethods)),
    ('tp_as_sequence', ctypes.POINTER(PySequenceMethods)),
    ('tp_as_mapping', ctypes.POINTER(PyMappingMethods)),
    ('tp_hash', ctypes.CFUNCTYPE(ctypes.c_int64, PyObject_p)),
    ('tp_call', ctypes.CFUNCTYPE(PyObject_p, PyObject_p, PyObject_p, PyObject_p)),
    ('tp_str', ctypes.CFUNCTYPE(PyObject_p, PyObject_p)),
    ('tp_getattro', ctypes.c_void_p),  # Type not declared yet
    ('tp_setattro', ctypes.c_void_p),  # Type not declared yet
    ('tp_as_buffer', ctypes.c_void_p),  # Type not declared yet
    ('tp_flags', ctypes.c_void_p),  # Type not declared yet
    ('tp_doc', ctypes.c_void_p),  # Type not declared yet
    ('tp_traverse', ctypes.c_void_p),  # Type not declared yet
    ('tp_clear', ctypes.c_void_p),  # Type not declared yet
    ('tp_richcompare', ctypes.c_void_p),  # Type not declared yet
    ('tp_weaklistoffset', ctypes.c_void_p),  # Type not declared yet
    ('tp_iter', ctypes.c_void_p),  # Type not declared yet
    ('tp_iternext', ctypes.c_void_p),  # Type not declared yet
    ('tp_methods', ctypes.c_void_p),  # Type not declared yet
    ('tp_members', ctypes.c_void_p),  # Type not declared yet
    ('tp_getset', ctypes.c_void_p),  # Type not declared yet
    ('tp_base', ctypes.c_void_p),  # Type not declared yet
    ('tp_dict', ctypes.c_void_p),  # Type not declared yet
    ('tp_descr_get', ctypes.c_void_p),  # Type not declared yet
    ('tp_descr_set', ctypes.c_void_p),  # Type not declared yet
    ('tp_dictoffset', Py_ssize_t),
    ('tp_init', ctypes.c_void_p),  # Type not declared yet
    ('tp_alloc', ctypes.c_void_p),  # Type not declared yet
    ('tp_new', ctypes.CFUNCTYPE(PyObject_p, PyObject_p, PyObject_p, ctypes.c_void_p)),
    ('tp_free', ctypes.CFUNCTYPE(None, ctypes.c_void_p)),
    ('tp_is_gc', Inquiry_p), # For PyObject_IS_GC
    ('tp_bases', PyObject_p),
    ('tp_mro', PyObject_p),
    ('tp_cache', PyObject_p),
    ('tp_subclasses', PyObject_p),
    ('tp_weaklist', PyObject_p),
    ('tp_del', Destructor_p),
    ('tp_version_tag', ctypes.c_uint),
    ('tp_finalize', Destructor_p),
] + ([
    ('tp_vectorcall', ctypes.CFUNCTYPE(PyObject_p, PyObject_p, ctypes.POINTER(PyObject_p), ctypes.c_size_t, PyObject_p)),
] if sys.hexversion >= 0x030800b1 else []) + ([
    ('tp_print', ctypes.c_void_p), # bpo-37250: kept for backwards compatibility in CPython 3.8 only
] if sys.hexversion >= 0x030800b2 and sys.hexversion < 0x03090000 else []) + ([ 
    ('tp_allocs', Py_ssize_t),
    ('tp_frees', Py_ssize_t),
    ('tp_maxalloc', Py_ssize_t),
    ('tp_prev', ctypes.POINTER(PyTypeObject)),
    ('tp_next', ctypes.POINTER(PyTypeObject)),
] if hasattr(sys, 'getcounts') else [])
    
assert ctypes.sizeof(PyTypeObject) == type.__sizeof__(object), "Size missmatch: defined {}, expected {}".format(ctypes.sizeof(PyTypeObject), type.__sizeof__(object))

# redundant dict of pointee types, because ctypes doesn't allow us
# to extract the pointee type from the pointer
PyTypeObject_as_types_dict = {
    'tp_as_async': PyAsyncMethods,
    'tp_as_number': PyNumberMethods,
    'tp_as_sequence': PySequenceMethods,
    'tp_as_mapping': PyMappingMethods,
}



def set_object_dict(obj, new_dict):
    from ctypes import (pythonapi, py_object, POINTER)
    from _ctypes import Py_INCREF as Py_IncRef
    from _ctypes import Py_DECREF as Py_DecRef 
#   documented, but much more complicated
#     Py_IncRef = pythonapi.Py_IncRef
#     Py_DecRef = pythonapi.Py_DecRef
#     Py_IncRef.argtypes = (py_object,); Py_IncRef.restype = None
#     Py_DecRef.argtypes = (py_object,); Py_DecRef.restype = None
    _PyObject_GetDictPtr = pythonapi._PyObject_GetDictPtr
    _PyObject_GetDictPtr.argtypes = (py_object,)
    _PyObject_GetDictPtr.restype = POINTER(py_object)

    try:
        from stackless import atomic
    except ImportError:
        from contextlib import nullcontext as atomic

    if not type(new_dict) is dict:
        raise TypeError('Argument new_dict must be a plain dictionary')
    try:
        old_dict = obj.__dict__
    except AttributeError as e:
        raise TypeError("object has no dictionary")
    if not type(old_dict) is dict:
        raise TypeError("object has no dictionary")
    
    if old_dict is new_dict:
        return old_dict

    with atomic():
        mod_dict = _PyObject_GetDictPtr(obj).contents
        assert mod_dict.value is old_dict
        Py_IncRef(new_dict)
        mod_dict.value = new_dict
        Py_DecRef(old_dict)
        mod_dict = None
    assert obj.__dict__ is new_dict
    return old_dict

if __name__ == '__main__':
    from types import ModuleType
    m = ModuleType('test')
    d = m.__dict__
    d_new = {}
    set_object_dict(m, d_new)
    assert m.__dict__ is d_new
   
    class C: pass
    set_object_dict(C(), d_new)