Support for Numpy BitGenerators PR#1 - Core Generator Support - #8031

from contextlib import ContextDecorator
import gc
import sys
import numpy as np

from numba import njit, types, objmode
from numba.extending import typeof_impl, models, register_model
from numba.extending import box, unbox, NativeValue
from numba.core import cgutils
from numba.core.runtime import rtsys


class Thing(object):
    def __init__(self):
        pass


class ThingType(types.Type):
    def __init__(self):
        super(ThingType, self).__init__(name='Thing')


thing_type = ThingType()


@typeof_impl.register(Thing)
def typeof_thing(val, c):
    return thing_type


@register_model(ThingType)
class ThingModel(models.StructModel):
    def __init__(self, dmm, fe_type):
        members = [
            ('meminfo', types.MemInfoPointer(types.voidptr)),
            ('parent', types.pyobject),
        ]
        models.StructModel.__init__(self, dmm, fe_type, members)


@unbox(ThingType)
def unbox_thing(typ, obj, c):
    thing_struct = cgutils.create_struct_proxy(typ)(c.context, c.builder)
    thing_struct.parent = obj
    NULL = cgutils.voidptr_t(None)
    thing_struct.meminfo = c.pyapi.nrt_meminfo_new_from_pyobject(
        NULL,  # there's no data
        obj,   # the python object, the call to nrt_meminfo_new_from_pyobject
               # will py_incref
    )
    is_py_error = cgutils.is_not_null(c.builder, c.pyapi.err_occurred())
    return NativeValue(thing_struct._getvalue(), is_error=is_py_error)


@box(ThingType)
def box_thing(typ, val, c):
    thing_struct = cgutils.create_struct_proxy(typ)(c.context, c.builder,
                                                    value=val)
    res = thing_struct.parent
    c.context.nrt.decref(c.builder, typ, val)
    return res


class memory_leak_check(ContextDecorator):

    def __init__(self, the_obj):
        self._the_obj = the_obj

    def __enter__(self):
        # Clean up any NRT-backed objects hanging in a dead reference cycle
        gc.collect()
        print(f"Python refcount: {sys.getrefcount(self._the_obj)}")
        self.__init_stats = rtsys.get_allocation_stats()

    def __exit__(self, *exc):
        gc.collect()
        old = self.__init_stats
        new = rtsys.get_allocation_stats()
        total_alloc = new.alloc - old.alloc
        total_free = new.free - old.free
        total_mi_alloc = new.mi_alloc - old.mi_alloc
        total_mi_free = new.mi_free - old.mi_free
        print(f"Allocs: {total_alloc}. Frees: {total_free}")
        print(f"MI Allocs: {total_mi_alloc}. MI Frees: {total_mi_free}")
        gc.collect()
        print(f"Python refcount: {sys.getrefcount(self._the_obj)}")


def np_ty(): return np.ones(3)


for Clazz in (np_ty, Thing):
    print('\n', f"Testing Type: {Clazz}".center(80, '='), '\n')
    print('Testing no-op'.center(80, '_'))

    thing_inst1 = Clazz()
    print(f"Python refcount: {sys.getrefcount(thing_inst1)}")  # should be 2
    with memory_leak_check(thing_inst1):
        thing_inst1  # Python refcount should be 3, the ctxmanager holds 1 more

    @njit
    def foo(x):
        return x

    thing_inst2 = Clazz()

    print('Testing foo'.center(80, '_'))
    with memory_leak_check(thing_inst2):
        initial_id = id(thing_inst2)
        result = foo(thing_inst2)
        result_id = id(result)
        print(f"ID check: {result is thing_inst2}")

    @njit
    def bar(x):
        pass  # x unused

    print('Testing bar'.center(80, '_'))
    thing_inst3 = Clazz()
    with memory_leak_check(thing_inst3):
        bar(thing_inst3)

    @njit
    def baz(x):
        with objmode():
            print(f"In objmode block {sys.getrefcount(x)}")

    print('Testing baz'.center(80, '_'))
    thing_inst4 = Clazz()
    with memory_leak_check(thing_inst4):
        baz(thing_inst4)