# pylint: disable=no-else-return, unidiomatic-typecheck, undefined-variable, wildcard-import """A global environment storing everything needed to interpret or compile a Relay program.""" from .base import register_relay_node, NodeBase from . import _make from . import _env @register_relay_node class Environment(NodeBase): """The global Relay environment containing functions, options and more. """ def __init__(self, funcs): """Construct an environment. Parameters ------ funcs: list of relay.Function Returns ------ env: A new environment containing :py:class:`~relay.env.Environment`. """ self.__init_handle_by_constructor__(_make.Environment, funcs) def add(self, var, func): """Add a function to the environment. Parameters --------- var: GlobalVar The global variable which names the function. func: Function The function. """ if isinstance(var, str): var = _env.Environment_GetGlobalVar(self, var) _env.Environment_Add(self, var, func) def merge(self, other): """Merge two environments. Parameters ---------- other: Environment The environment to merge into the current Environment. """ return _env.Environment_Merge(self, other) def global_var(self, name): """Get a global variable by name. Parameters ---------- name: str The name of the global variable. Returns ------- global_var: GlobalVar The global variable mapped to :code:`name`. """ return _env.Environment_GetGlobalVar(self, name) def __getitem__(self, var): """Lookup a global function by name or by variable. Parameters ---------- var: str or GlobalVar The name or global variable. Returns ------- func: Function The function referenced by :code:`var`. """ if isinstance(var, str): return _env.Environment_Lookup_str(self, var) else: return _env.Environment_Lookup(self, var)