funcdesc.py

"""
Function descriptors.
"""

from collections import defaultdict
import importlib

from numba.core import types, itanium_mangler
from numba.core.utils import _dynamic_modname, _dynamic_module


def default_mangler(name, argtypes, *, abi_tags=(), uid=None):
    return itanium_mangler.mangle(name, argtypes, abi_tags=abi_tags, uid=uid)


def qualifying_prefix(modname, qualname):
    """
    Returns a new string that is used for the first half of the mangled name.
    """
    # XXX choose a different convention for object mode
    return '{}.{}'.format(modname, qualname) if modname else qualname


class FunctionDescriptor(object):
    """
    Base class for function descriptors: an object used to carry
    useful metadata about a natively callable function.

    Note that while `FunctionIdentity` denotes a Python function
    which is being concretely compiled by Numba, `FunctionDescriptor`
    may be more "abstract": e.g. a function decorated with `@generated_jit`.
    """
    __slots__ = ('native', 'modname', 'qualname', 'doc', 'typemap',
                 'calltypes', 'args', 'kws', 'restype', 'argtypes',
                 'mangled_name', 'unique_name', 'env_name', 'global_dict',
                 'inline', 'noalias', 'abi_tags', 'uid')

    def __init__(self, native, modname, qualname, unique_name, doc,
                 typemap, restype, calltypes, args, kws, mangler=None,
                 argtypes=None, inline=False, noalias=False, env_name=None,
                 global_dict=None, abi_tags=(), uid=None):
        self.native = native
        self.modname = modname
        self.global_dict = global_dict
        self.qualname = qualname
        self.unique_name = unique_name
        self.doc = doc
        # XXX typemap and calltypes should be on the compile result,
        # not the FunctionDescriptor
        self.typemap = typemap
        self.calltypes = calltypes
        self.args = args
        self.kws = kws
        self.restype = restype
        # Argument types
        if argtypes is not None:
            assert isinstance(argtypes, tuple), argtypes
            self.argtypes = argtypes
        else:
            # Get argument types from the type inference result
            # (note the "arg.FOO" convention as used in typeinfer
            self.argtypes = tuple(self.typemap['arg.' + a] for a in args)
        mangler = default_mangler if mangler is None else mangler
        # The mangled name *must* be unique, else the wrong function can
        # be chosen at link time.
        qualprefix = qualifying_prefix(self.modname, self.qualname)
        self.uid = uid
        self.mangled_name = mangler(
            qualprefix, self.argtypes, abi_tags=abi_tags, uid=uid,
        )
        if env_name is None:
            env_name = mangler(".NumbaEnv.{}".format(qualprefix),
                               self.argtypes, abi_tags=abi_tags, uid=uid)
        self.env_name = env_name
        self.inline = inline
        self.noalias = noalias
        self.abi_tags = abi_tags

    def lookup_globals(self):
        """
        Return the global dictionary of the function.
        It may not match the Module's globals if the function is created
        dynamically (i.e. exec)
        """
        return self.global_dict or self.lookup_module().__dict__

    def lookup_module(self):
        """
        Return the module in which this function is supposed to exist.
        This may be a dummy module if the function was dynamically
        generated or the module can't be found.
        """
        if self.modname == _dynamic_modname:
            return _dynamic_module
        else:
            try:
                # ensure module exist
                return importlib.import_module(self.modname)
            except ImportError:
                return _dynamic_module

    def lookup_function(self):
        """
        Return the original function object described by this object.
        """
        return getattr(self.lookup_module(), self.qualname)

    @property
    def llvm_func_name(self):
        """
        The LLVM-registered name for the raw function.
        """
        return self.mangled_name

    # XXX refactor this

    @property
    def llvm_cpython_wrapper_name(self):
        """
        The LLVM-registered name for a CPython-compatible wrapper of the
        raw function (i.e. a PyCFunctionWithKeywords).
        """
        return itanium_mangler.prepend_namespace(self.mangled_name,
                                                 ns='cpython')

    @property
    def llvm_cfunc_wrapper_name(self):
        """
        The LLVM-registered name for a C-compatible wrapper of the
        raw function.
        """
        return 'cfunc.' + self.mangled_name

    def __repr__(self):
        return "<function descriptor %r>" % (self.unique_name)

    @classmethod
    def _get_function_info(cls, func_ir):
        """
        Returns
        -------
        qualname, unique_name, modname, doc, args, kws, globals

        ``unique_name`` must be a unique name.
        """
        func = func_ir.func_id.func
        qualname = func_ir.func_id.func_qualname
        # XXX to func_id
        modname = func.__module__
        doc = func.__doc__ or ''
        args = tuple(func_ir.arg_names)
        kws = ()        # TODO
        global_dict = None

        if modname is None:
            # Dynamically generated function.
            modname = _dynamic_modname
            # Retain a reference to the dictionary of the function.
            # This disables caching, serialization and pickling.
            global_dict = func_ir.func_id.func.__globals__

        unique_name = func_ir.func_id.unique_name

        return qualname, unique_name, modname, doc, args, kws, global_dict

    @classmethod
    def _from_python_function(cls, func_ir, typemap, restype,
                              calltypes, native, mangler=None,
                              inline=False, noalias=False, abi_tags=()):
        (qualname, unique_name, modname, doc, args, kws, global_dict,
         ) = cls._get_function_info(func_ir)

        self = cls(native, modname, qualname, unique_name, doc,
                   typemap, restype, calltypes,
                   args, kws, mangler=mangler, inline=inline, noalias=noalias,
                   global_dict=global_dict, abi_tags=abi_tags,
                   uid=func_ir.func_id.unique_id)
        return self


class PythonFunctionDescriptor(FunctionDescriptor):
    """
    A FunctionDescriptor subclass for Numba-compiled functions.
    """
    __slots__ = ()

    @classmethod
    def from_specialized_function(cls, func_ir, typemap, restype, calltypes,
                                  mangler, inline, noalias, abi_tags):
        """
        Build a FunctionDescriptor for a given specialization of a Python
        function (in nopython mode).
        """
        return cls._from_python_function(func_ir, typemap, restype, calltypes,
                                         native=True, mangler=mangler,
                                         inline=inline, noalias=noalias,
                                         abi_tags=abi_tags)

    @classmethod
    def from_object_mode_function(cls, func_ir):
        """
        Build a FunctionDescriptor for an object mode variant of a Python
        function.
        """
        typemap = defaultdict(lambda: types.pyobject)
        calltypes = typemap.copy()
        restype = types.pyobject
        return cls._from_python_function(func_ir, typemap, restype, calltypes,
                                         native=False)


class ExternalFunctionDescriptor(FunctionDescriptor):
    """
    A FunctionDescriptor subclass for opaque external functions
    (e.g. raw C functions).
    """
    __slots__ = ()

    def __init__(self, name, restype, argtypes):
        args = ["arg%d" % i for i in range(len(argtypes))]

        def mangler(a, x, abi_tags, uid=None):
            return a
        super(ExternalFunctionDescriptor, self
              ).__init__(native=True, modname=None, qualname=name,
                         unique_name=name, doc='', typemap=None,
                         restype=restype, calltypes=None, args=args,
                         kws=None,
                         mangler=mangler,
                         argtypes=argtypes)