<commit_before>from __future__ import print_function, division, absolute_import
from pprint import pprint
from contextlib import contextmanager
from collections import namedtuple, defaultdict
from numba import (bytecode, interpreter, typing, typeinfer, lowering,
                   irpasses, utils, config, type_annotations, types, ir,
                   assume)
from numba.targets import cpu


class Flags(utils.ConfigOptions):
    OPTIONS = frozenset(['enable_pyobject',
                         'force_pyobject',
                         'no_compile'])


DEFAULT_FLAGS = Flags()


CR_FIELDS = ["typing_context",
             "target_context",
             "entry_point",
             "entry_point_addr",
             "typing_error",
             "type_annotation",
             "llvm_module",
             "llvm_func",
             "signature",
             "objectmode",]


CompileResult = namedtuple("CompileResult", CR_FIELDS)


def compile_result(**kws):
    keys = set(kws.keys())
    fieldset = set(CR_FIELDS)
    badnames = keys - fieldset
    if badnames:
        raise NameError(*badnames)
    missing = fieldset - keys
    for k in missing:
        kws[k] = None
    return CompileResult(**kws)


def compile_isolated(func, args, return_type=None, flags=DEFAULT_FLAGS,
                     locals={}):
    """
    Compile the function is an isolated environment.
    Good for testing.
    """
    typingctx = typing.Context()
    targetctx = cpu.CPUContext(typingctx)
    return compile_extra(typingctx, targetctx, func, args, return_type, flags,
                         locals)


class _CompileStatus(object):
    """
    Used like a C record
    """
    __slots__ = 'fail_reason', 'use_python_mode', 'can_fallback'


@contextmanager
def _fallback_context(status):
    try:
        yield
    except Exception as e:
        if not status.can_fallback:
            raise
        status.fail_reason = e
        status.use_python_mode = True


def compile_extra(typingctx, targetctx, func, args, return_type, flags,
                  locals):
    """
    Args
    ----
    - return_type
        Use ``None`` to indicate
    """
    # Translate to IR
    interp = translate_stage(func)
    nargs = len(interp.argspec.args)

    status = _CompileStatus()
    status.can_fallback = flags.enable_pyobject
    status.fail_reason = None
    status.use_python_mode = flags.force_pyobject

    localctx = targetctx.localized()

    if not status.use_python_mode:
        with _fallback_context(status):
            # Type inference
            typemap, return_type, calltypes = type_inference_stage(typingctx,
                                                                   interp,
                                                                   args,
                                                                   return_type,
                                                                   locals)

        if not status.use_python_mode:
            with _fallback_context(status):
                legalize_return_type(return_type, interp, localctx)

    if status.use_python_mode:
        func, fnptr, lmod, lfunc = py_lowering_stage(localctx, interp,
                                                     flags.no_compile)
        typemap = defaultdict(lambda: types.pyobject)
        calltypes = defaultdict(lambda: types.pyobject)

        return_type = types.pyobject
        args = [types.pyobject] * nargs
    else:
        func, fnptr, lmod, lfunc = native_lowering_stage(localctx, interp,
                                                         typemap,
                                                         return_type,
                                                         calltypes,
                                                         flags.no_compile)

    type_annotation = type_annotations.TypeAnnotation(interp=interp,
                                                      typemap=typemap,
                                                      calltypes=calltypes)

    signature = typing.signature(return_type, *args)

    assert lfunc.module is lmod
    cr = compile_result(typing_context=typingctx,
                        target_context=targetctx,
                        entry_point=func,
                        entry_point_addr=fnptr,
                        typing_error=status.fail_reason,
                        type_annotation=type_annotation,
                        llvm_func=lfunc,
                        llvm_module=lmod,
                        signature=signature,
                        objectmode=status.use_python_mode)
    return cr


def legalize_return_type(return_type, interp, targetctx):
    """
    Only accept array return type iff it is passed into the function.<commit_msg>Prevent pyobject types in nopython mode<commit_after>from __future__ import print_function, division, absolute_import
from pprint import pprint
from contextlib import contextmanager
from collections import namedtuple, defaultdict
from numba import (bytecode, interpreter, typing, typeinfer, lowering,
                   irpasses, utils, config, type_annotations, types, ir,
                   assume)
from numba.targets import cpu


class Flags(utils.ConfigOptions):
    OPTIONS = frozenset(['enable_pyobject',
                         'force_pyobject',
                         'no_compile'])


DEFAULT_FLAGS = Flags()


CR_FIELDS = ["typing_context",
             "target_context",
             "entry_point",
             "entry_point_addr",
             "typing_error",
             "type_annotation",
             "llvm_module",
             "llvm_func",
             "signature",
             "objectmode",]


CompileResult = namedtuple("CompileResult", CR_FIELDS)


def compile_result(**kws):
    keys = set(kws.keys())
    fieldset = set(CR_FIELDS)
    badnames = keys - fieldset
    if badnames:
        raise NameError(*badnames)
    missing = fieldset - keys
    for k in missing:
        kws[k] = None
    return CompileResult(**kws)


def compile_isolated(func, args, return_type=None, flags=DEFAULT_FLAGS,
                     locals={}):
    """
    Compile the function is an isolated environment.
    Good for testing.
    """
    typingctx = typing.Context()
    targetctx = cpu.CPUContext(typingctx)
    return compile_extra(typingctx, targetctx, func, args, return_type, flags,
                         locals)


class _CompileStatus(object):
    """
    Used like a C record
    """
    __slots__ = 'fail_reason', 'use_python_mode', 'can_fallback'


@contextmanager
def _fallback_context(status):
    try:
        yield
    except Exception as e:
        if not status.can_fallback:
            raise
        status.fail_reason = e
        status.use_python_mode = True


def compile_extra(typingctx, targetctx, func, args, return_type, flags,
                  locals):
    """
    Args
    ----
    - return_type
        Use ``None`` to indicate
    """
    # Translate to IR
    interp = translate_stage(func)
    nargs = len(interp.argspec.args)

    status = _CompileStatus()
    status.can_fallback = flags.enable_pyobject
    status.fail_reason = None
    status.use_python_mode = flags.force_pyobject

    localctx = targetctx.localized()

    if not status.use_python_mode:
        with _fallback_context(status):
            legialize_given_types(args, return_type)
            # Type inference
            typemap, return_type, calltypes = type_inference_stage(typingctx,
                                                                   interp,
                                                                   args,
                                                                   return_type,
                                                                   locals)

        if not status.use_python_mode:
            with _fallback_context(status):
                legalize_return_type(return_type, interp, localctx)

    if status.use_python_mode:
        func, fnptr, lmod, lfunc = py_lowering_stage(localctx, interp,
                                                     flags.no_compile)
        typemap = defaultdict(lambda: types.pyobject)
        calltypes = defaultdict(lambda: types.pyobject)

        return_type = types.pyobject
        args = [types.pyobject] * nargs
    else:
        func, fnptr, lmod, lfunc = native_lowering_stage(localctx, interp,
                                                         typemap,
                                                         return_type,
                                                         calltypes,
                                                         flags.no_compile)

    type_annotation = type_annotations.TypeAnnotation(interp=interp,
                                                      typemap=typemap,
                                                      calltypes=calltypes)

    signature = typing.signature(return_type, *args)

    assert lfunc.module is lmod
    cr = compile_result(typing_context=typingctx,
                        target_context=targetctx,
                        entry_point=func,
                        entry_point_addr=fnptr,
                        typing_error=status.fail_reason,
                        type_annotation=type_annotation,
                        llvm_func=lfunc,
                        llvm_module=lmod,
                        signature=signature,
                        objectmode=status.use_python_mode)
    return cr


def _is_nopython_types(t):
    return t != types.pyobject and not isinstance(t, types.Dummy)


def legialize_given_types(args, return_type):
    # Filter argument types
    for i, a in enumerate(args):
        if not _is_nopython_types(a):
            raise TypeError("Arg %d of %s is not legal in nopython "
                            "mode" % (i, a))
    # Filter return type
    if (return_type and return_type != types.none and
            not _is_nopython_types(return_type)):
        raise TypeError('Return type of %s is not legal in nopython '
                        'mode' % (return_type,))


def legalize_return_type(return_type, interp, targetctx):
    """
    Only accept array return type iff it is passed into the function.
{'fraction_non_alphanumeric': 0.07177221404025527, 'fraction_numerical': 0.0, 'mean_word_length': 1.5683308119011599, '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': 44, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5677488', 'n_tokens_mistral': 2619, 'n_tokens_neox': 2522, 'n_words': 595}