# pylint: disable=invalid-name, unused-import """Function namespace.""" from __future__ import absolute_import import ctypes import sys from .base import _FFI_MODE, _LIB, c_str, check_call, py_str, string_types IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError try: # pylint: disable=wrong-import-position if _FFI_MODE == "ctypes": raise ImportError() if sys.version_info >= (3, 0): from ._cy3.core import ( _set_class_function, _set_class_module, convert_to_dgl_func, FunctionBase as _FunctionBase, ) else: from ._cy2.core import ( _set_class_function, _set_class_module, convert_to_dgl_func, FunctionBase as _FunctionBase, ) except IMPORT_EXCEPT: # pylint: disable=wrong-import-position from ._ctypes.function import ( _set_class_function, _set_class_module, convert_to_dgl_func, FunctionBase as _FunctionBase, ) FunctionHandle = ctypes.c_void_p class Function(_FunctionBase): """The PackedFunc object. Function plays an key role to bridge front and backend in DGL. Function provide a type-erased interface, you can call function with positional arguments. The compiled module returns Function. DGL backend also registers and exposes its API as Functions. For example, the developer function exposed in dgl.ir_pass are actually C++ functions that are registered as PackedFunc The following are list of common usage scenario of dgl.Function. - Automatic exposure of C++ API into python - To call PackedFunc from python side - To call python callbacks to inspect results in generated code - Bring python hook into C++ backend See Also -------- dgl.register_func: How to register global function. dgl.get_global_func: How to get global function. """ pass # pylint: disable=unnecessary-pass class ModuleBase(object): """Base class for module""" __slots__ = ["handle", "_entry", "entry_name"] def __init__(self, handle): self.handle = handle self._entry = None self.entry_name = "__dgl_main__" def __del__(self): check_call(_LIB.DGLModFree(self.handle)) @property def entry_func(self): """Get the entry function Returns ------- f : Function The entry function if exist """ if self._entry: return self._entry self._entry = self.get_function(self.entry_name) return self._entry def get_function(self, name, query_imports=False): """Get function from the module. Parameters ---------- name : str The name of the function query_imports : bool Whether also query modules imported by this module. Returns ------- f : Function The result function. """ ret_handle = FunctionHandle() check_call( _LIB.DGLModGetFunction( self.handle, c_str(name), ctypes.c_int(query_imports), ctypes.byref(ret_handle), ) ) if not ret_handle.value: raise AttributeError("Module has no function '%s'" % name) return Function(ret_handle, False) def import_module(self, module): """Add module to the import list of current one. Parameters ---------- module : Module The other module. """ check_call(_LIB.DGLModImport(self.handle, module.handle)) def __getitem__(self, name): if not isinstance(name, string_types): raise ValueError("Can only take string as function name") return self.get_function(name) def __call__(self, *args): if self._entry: return self._entry(*args) f = self.entry_func return f(*args) def register_func(func_name, f=None, override=False): """Register global function Parameters ---------- func_name : str or function The function name f : function, optional The function to be registered. override: boolean optional Whether override existing entry. Returns ------- fregister : function Register function if f is not specified. Examples -------- The following code registers my_packed_func as global function. Note that we simply get it back from global function table to invoke it from python side. However, we can also invoke the same function from C++ backend, or in the compiled DGL code. .. code-block:: python targs = (10, 10.0, "hello") @dgl.register_func def my_packed_func(*args): assert(tuple(args) == targs) return 10 # Get it out from global function table f = dgl.get_global_func("my_packed_func") assert isinstance(f, dgl.nd.Function) y = f(*targs) assert y == 10 """ if callable(func_name): f = func_name func_name = f.__name__ if not isinstance(func_name, str): raise ValueError("expect string function name") ioverride = ctypes.c_int(override) def register(myf): """internal register function""" if not isinstance(myf, Function): myf = convert_to_dgl_func(myf) check_call( _LIB.DGLFuncRegisterGlobal(c_str(func_name), myf.handle, ioverride) ) return myf if f: return register(f) return register def get_global_func(name, allow_missing=False): """Get a global function by name Parameters ---------- name : str The name of the global function allow_missing : bool Whether allow missing function or raise an error. Returns ------- func : dgl.Function The function to be returned, None if function is missing. """ handle = FunctionHandle() check_call(_LIB.DGLFuncGetGlobal(c_str(name), ctypes.byref(handle))) if handle.value: return Function(handle, False) else: if allow_missing: return None else: raise ValueError("Cannot find global function %s" % name) def list_global_func_names(): """Get list of global functions registered. Returns ------- names : list List of global functions names. """ plist = ctypes.POINTER(ctypes.c_char_p)() size = ctypes.c_uint() check_call( _LIB.DGLFuncListGlobalNames(ctypes.byref(size), ctypes.byref(plist)) ) fnames = [] for i in range(size.value): fnames.append(py_str(plist[i])) return fnames def extract_ext_funcs(finit): """ Extract the extension PackedFuncs from a C module. Parameters ---------- finit : ctypes function a ctypes that takes signature of DGLExtensionDeclarer Returns ------- fdict : dict of str to Function The extracted functions """ fdict = {} def _list(name, func): fdict[name] = func myf = convert_to_dgl_func(_list) ret = finit(myf.handle) _ = myf if ret != 0: raise RuntimeError("cannot initialize with %s" % finit) return fdict def _get_api(f): flocal = f flocal.is_global = True return flocal def _init_api(namespace, target_module_name=None): """Initialize api for a given module name namespace : str The namespace of the source registry target_module_name : str The target module name if different from namespace """ target_module_name = target_module_name if target_module_name else namespace if namespace.startswith("dgl."): return _init_api_prefix(target_module_name, namespace[4:]) else: return _init_api_prefix(target_module_name, namespace) def _init_api_prefix(module_name, prefix): module = sys.modules[module_name] name_list = [] for name in list_global_func_names(): if name.startswith("_") and not name.startswith("_deprecate"): # internal APIs are ignored continue name_split = name.rsplit(".", 1) if name_split[0] != prefix: continue if len(name_split) == 1: print('Warning: invalid API name "%s".' % name) continue fname = name_split[1] target_module = module f = get_global_func(name) ff = _get_api(f) ff.__name__ = fname ff.__doc__ = "DGL PackedFunc %s. " % fname setattr(target_module, ff.__name__, ff) name_list.append(fname) return name_list def _init_internal_api(): for name in list_global_func_names(): if not name.startswith("_") or name.startswith("_deprecate"): # normal APIs are ignored continue target_module = sys.modules["dgl._api_internal"] fname = name if fname.find(".") != -1: print('Warning: invalid API name "%s".' % fname) continue f = get_global_func(name) ff = _get_api(f) ff.__name__ = fname ff.__doc__ = "DGL PackedFunc %s. " % fname setattr(target_module, ff.__name__, ff) _set_class_function(Function)