<commit_before># -*- coding: utf-8 -*-
from __future__ import print_function, division, absolute_import
import warnings
import numpy as np
from numba.decorators import jit
from numba.targets.registry import target_registry
from numba.targets.descriptors import TargetDescriptor
from numba.targets.options import TargetOptions
from numba import utils, compiler, types, sigutils
from . import _internal
from .sigparse import parse_signature
from .wrappers import build_ufunc_wrapper, build_gufunc_wrapper
from numba.targets import registry


class UFuncTargetOptions(TargetOptions):
    OPTIONS = {
        "nopython" : bool,
        "forceobj" : bool,
    }


class UFuncTarget(registry.CPUTarget):
    options = UFuncTargetOptions


class UFuncDispatcher(object):
    targetdescr = UFuncTarget()

    def __init__(self, py_func, locals={}, targetoptions={}):
        self.py_func = py_func
        self.overloads = utils.UniqueDict()
        self.targetoptions = targetoptions
        self.locals = locals

    def compile(self, sig, locals={}, **targetoptions):
        locs = self.locals.copy()
        locs.update(locals)

        topt = self.targetoptions.copy()
        topt.update(targetoptions)

        flags = compiler.Flags()
        self.targetdescr.options.parse_as_flags(flags, topt)
        flags.set("no_compile")
        # Disable loop lifting
        # The feature requires a real python function
        flags.unset("enable_looplift")

        typingctx = self.targetdescr.typing_context
        targetctx = self.targetdescr.target_context

        args, return_type = sigutils.normalize_signature(sig)

        cres = compiler.compile_extra(typingctx, targetctx, self.py_func,
                                      args=args, return_type=return_type,
                                      flags=flags, locals=locals)

        self.overloads[cres.signature] = cres
        return cres


target_registry['npyufunc'] = UFuncDispatcher


class UFuncBuilder(object):
    def __init__(self, py_func, targetoptions={}):
        self.py_func = py_func
        self.nb_func = jit(target='npyufunc', **targetoptions)(py_func)

    def add(self, sig=None, argtypes=None, restype=None):
        # Handle argtypes
        if argtypes is not None:
            warnings.warn("Keyword argument argtypes is deprecated",
                          DeprecationWarning)
            assert sig is None
            if restype is None:
                sig = tuple(argtypes)
            else:
                sig = restype(*argtypes)
        # Do compilation
        # Return CompileResult to test
        return self.nb_func.compile(sig)

    def build_ufunc(self):
        dtypelist = []
        ptrlist = []
        if not self.nb_func:
            raise TypeError("No definition")
        for sig, cres in self.nb_func.overloads.items():
            dtypenums, ptr = self.build(cres)
            dtypelist.append(dtypenums)
            ptrlist.append(utils.longint(ptr))
        datlist = [None] * len(ptrlist)

        inct = len(cres.signature.args)
        outct = 1

        # Becareful that fromfunc does not provide full error checking yet.
        # If typenum is out-of-bound, we have nasty memory corruptions.
        # For instance, -1 for typenum will cause segfault.
        # If elements of type-list (2nd arg) is tuple instead,
        # there will also memory corruption. (Seems like code rewrite.)
        ufunc = _internal.fromfunc(ptrlist, dtypelist, inct, outct, datlist)

        return ufunc

    def build(self, cres):
        # Buider wrapper for ufunc entry point
        ctx = cres.target_context
        signature = cres.signature
        wrapper = build_ufunc_wrapper(ctx, cres.llvm_func, signature)
        ctx.finalize(wrapper, cres.fndesc)
        ctx.engine.add_module(wrapper.module)
        ptr = ctx.engine.get_pointer_to_function(wrapper)
        # Get dtypes
        dtypenums = [np.dtype(a.name).num for a in signature.args]
        dtypenums.append(np.dtype(signature.return_type.name).num)
        return dtypenums, ptr


class GUFuncBuilder(object):
    # TODO handle scalar
    def __init__(self, py_func, signature, targetoptions={}):
        self.py_func = py_func
        self.nb_func = jit(target='npyufunc')(py_func)
        self.signature = signature
        self.sin, self.sout = parse_signature(signature)
        self.targetoptions = targetoptions
        self._envs = []

    def add(self, sig=None, argtypes=None, restype=None):
        # Handle argtypes
        if argtypes is not None:
            warnings.warn("Keyword argument argtypes is deprecated",
                          DeprecationWarning)
            assert sig is None
            if restype is None:
                sig = tuple(argtypes)
            else:
                sig = restype(*argtypes)
        # Actual work begins
        cres = self.nb_func.compile(sig, **self.targetoptions)

        if not cres.objectmode and cres.signature.return_type != types.void:
            raise TypeError("gufunc kernel must have void return type")

        return cres

    def build_ufunc(self):
        dtypelist = []
        ptrlist = []
        if not self.nb_func:
            raise TypeError("No definition")

        for sig, cres in self.nb_func.overloads.items():
            dtypenums, ptr, env = self.build(cres)
            self._envs.append(env)   # keep env object alive
            dtypelist.append(dtypenums)
            ptrlist.append(utils.longint(ptr))
        datlist = [None] * len(ptrlist)

        inct = len(self.sin)
        outct = len(self.sout)

        ufunc = _internal.fromfuncsig(ptrlist, dtypelist, inct, outct, datlist,
                                      self.signature)

        return ufunc

    def build(self, cres):
        """
        Returns (dtype numbers, function ptr, EnviornmentObject)
        """
        # Buider wrapper for ufunc entry point
        ctx = cres.target_context
        signature = cres.signature
        wrapper, env = build_gufunc_wrapper(ctx, cres.llvm_func, signature,
                                            self.sin, self.sout,
                                            fndesc=cres.fndesc)

        ctx.finalize(wrapper, cres.fndesc)
        ctx.engine.add_module(wrapper.module)
        ptr = ctx.engine.get_pointer_to_function(wrapper)
        # Get dtypes
        dtypenums = []
        for a in signature.args:
            if isinstance(a, types.Array):
                ty = a.dtype
            else:
                ty = a
            dtypenums.append(np.dtype(ty.name).num)
        return dtypenums, ptr, env

<commit_msg>Fix env keepalive<commit_after># -*- coding: utf-8 -*-
from __future__ import print_function, division, absolute_import
import warnings
import numpy as np
from numba.decorators import jit
from numba.targets.registry import target_registry
from numba.targets.descriptors import TargetDescriptor
from numba.targets.options import TargetOptions
from numba import utils, compiler, types, sigutils
from . import _internal
from .sigparse import parse_signature
from .wrappers import build_ufunc_wrapper, build_gufunc_wrapper
from numba.targets import registry


class UFuncTargetOptions(TargetOptions):
    OPTIONS = {
        "nopython" : bool,
        "forceobj" : bool,
    }


class UFuncTarget(registry.CPUTarget):
    options = UFuncTargetOptions


class UFuncDispatcher(object):
    targetdescr = UFuncTarget()

    def __init__(self, py_func, locals={}, targetoptions={}):
        self.py_func = py_func
        self.overloads = utils.UniqueDict()
        self.targetoptions = targetoptions
        self.locals = locals

    def compile(self, sig, locals={}, **targetoptions):
        locs = self.locals.copy()
        locs.update(locals)

        topt = self.targetoptions.copy()
        topt.update(targetoptions)

        flags = compiler.Flags()
        self.targetdescr.options.parse_as_flags(flags, topt)
        flags.set("no_compile")
        # Disable loop lifting
        # The feature requires a real python function
        flags.unset("enable_looplift")

        typingctx = self.targetdescr.typing_context
        targetctx = self.targetdescr.target_context

        args, return_type = sigutils.normalize_signature(sig)

        cres = compiler.compile_extra(typingctx, targetctx, self.py_func,
                                      args=args, return_type=return_type,
                                      flags=flags, locals=locals)

        self.overloads[cres.signature] = cres
        return cres


target_registry['npyufunc'] = UFuncDispatcher


class UFuncBuilder(object):
    def __init__(self, py_func, targetoptions={}):
        self.py_func = py_func
        self.nb_func = jit(target='npyufunc', **targetoptions)(py_func)

    def add(self, sig=None, argtypes=None, restype=None):
        # Handle argtypes
        if argtypes is not None:
            warnings.warn("Keyword argument argtypes is deprecated",
                          DeprecationWarning)
            assert sig is None
            if restype is None:
                sig = tuple(argtypes)
            else:
                sig = restype(*argtypes)
        # Do compilation
        # Return CompileResult to test
        return self.nb_func.compile(sig)

    def build_ufunc(self):
        dtypelist = []
        ptrlist = []
        if not self.nb_func:
            raise TypeError("No definition")
        for sig, cres in self.nb_func.overloads.items():
            dtypenums, ptr = self.build(cres)
            dtypelist.append(dtypenums)
            ptrlist.append(utils.longint(ptr))
        datlist = [None] * len(ptrlist)

        inct = len(cres.signature.args)
        outct = 1

        # Becareful that fromfunc does not provide full error checking yet.
        # If typenum is out-of-bound, we have nasty memory corruptions.
        # For instance, -1 for typenum will cause segfault.
        # If elements of type-list (2nd arg) is tuple instead,
        # there will also memory corruption. (Seems like code rewrite.)
        ufunc = _internal.fromfunc(ptrlist, dtypelist, inct, outct, datlist)

        return ufunc

    def build(self, cres):
        # Buider wrapper for ufunc entry point
        ctx = cres.target_context
        signature = cres.signature
        wrapper = build_ufunc_wrapper(ctx, cres.llvm_func, signature)
        ctx.finalize(wrapper, cres.fndesc)
        ctx.engine.add_module(wrapper.module)
        ptr = ctx.engine.get_pointer_to_function(wrapper)
        # Get dtypes
        dtypenums = [np.dtype(a.name).num for a in signature.args]
        dtypenums.append(np.dtype(signature.return_type.name).num)
        return dtypenums, ptr


class GUFuncBuilder(object):
    # TODO handle scalar
    def __init__(self, py_func, signature, targetoptions={}):
        self.py_func = py_func
        self.nb_func = jit(target='npyufunc')(py_func)
        self.signature = signature
        self.sin, self.sout = parse_signature(signature)
        self.targetoptions = targetoptions

    def add(self, sig=None, argtypes=None, restype=None):
        # Handle argtypes
        if argtypes is not None:
            warnings.warn("Keyword argument argtypes is deprecated",
                          DeprecationWarning)
            assert sig is None
            if restype is None:
                sig = tuple(argtypes)
            else:
                sig = restype(*argtypes)
        # Actual work begins
        cres = self.nb_func.compile(sig, **self.targetoptions)

        if not cres.objectmode and cres.signature.return_type != types.void:
            raise TypeError("gufunc kernel must have void return type")

        return cres

    def build_ufunc(self):
        dtypelist = []
        ptrlist = []
        keepalive = []

        if not self.nb_func:
            raise TypeError("No definition")

        for sig, cres in self.nb_func.overloads.items():
            dtypenums, ptr, env = self.build(cres)
            keepalive.append(env)   # keep env object alive
            dtypelist.append(dtypenums)
            ptrlist.append(utils.longint(ptr))
        datlist = [None] * len(ptrlist)

        inct = len(self.sin)
        outct = len(self.sout)

        # Pass envs to fromfuncsig to bind to the lifetime of the ufunc object
        ufunc = _internal.fromfuncsig(ptrlist, dtypelist, inct, outct, datlist,
                                      self.signature, keepalive)
        return ufunc

    def build(self, cres):
        """
        Returns (dtype numbers, function ptr, EnviornmentObject)
        """
        # Buider wrapper for ufunc entry point
        ctx = cres.target_context
        signature = cres.signature
        wrapper, env = build_gufunc_wrapper(ctx, cres.llvm_func, signature,
                                            self.sin, self.sout,
                                            fndesc=cres.fndesc)

        ctx.finalize(wrapper, cres.fndesc)
        ctx.engine.add_module(wrapper.module)
        ptr = ctx.engine.get_pointer_to_function(wrapper)
        # Get dtypes
        dtypenums = []
        for a in signature.args:
            if isinstance(a, types.Array):
                ty = a.dtype
            else:
                ty = a
            dtypenums.append(np.dtype(ty.name).num)
        return dtypenums, ptr, env


{'fraction_non_alphanumeric': 0.08559843820393452, 'fraction_numerical': 0.0006006907944135756, 'mean_word_length': 2.1486997635933807, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5977318', 'n_tokens_mistral': 3862, 'n_tokens_neox': 3729, 'n_words': 942}