diff --git a/src/celpy/evaluation.py b/src/celpy/evaluation.py index e311801..b2ad11d 100644 --- a/src/celpy/evaluation.py +++ b/src/celpy/evaluation.py @@ -3059,6 +3059,9 @@ def __init__( self.base_activation = activation self.activation = self.base_activation + # Top-level modules for the functions, so ``module.function`` names resolve. + self.fn_packages = self._resolve_fn_packages() + self.logger.debug("Transpiler activation: %r", self.activation) # self.logger.debug("functions: %r", self.functions) # Refactor ``self.functions`` into an Activation @@ -3089,8 +3092,9 @@ def evaluate(self, context: Context) -> celpy.celtypes.Value: self.logger.debug("Activation: %r", self.activation) # Global for the top-level ``CEL = result(base_activation, ...)`` statement. + # Function packages (e.g., ``__main__``) are added; module names take precedence. evaluation_globals = ( - celpy.evaluation.result.__globals__ + self.fn_packages | celpy.evaluation.result.__globals__ ) # the ``evaluation`` moodule evaluation_globals["base_activation"] = self.activation try: @@ -3104,6 +3108,23 @@ def evaluate(self, context: Context) -> celpy.celtypes.Value: self.logger.error("Internal error: %r", ex) raise CELEvalError("evaluation error", type(ex), ex.args) + def _resolve_fn_packages(self) -> dict[str, Any]: + """ + Locate the top-level module for each function in the activation. + + The transpiled code refers to a function as ``module.qualname``. + The top-level package of that module must be a global name during :py:meth:`evaluate`. + """ + packages: dict[str, Any] = {} + for func in self.base_activation.functions.values(): + module_name = {"_operator": "operator"}.get( + func.__module__, func.__module__ + ) + package = module_name.split(".", 1)[0] + if package not in packages and package in sys.modules: + packages[package] = sys.modules[package] + return packages + class Phase1Transpiler(lark.visitors.Visitor_Recursive): """ diff --git a/tests/test_transpilation.py b/tests/test_transpilation.py index 176c9dc..48acd5f 100644 --- a/tests/test_transpilation.py +++ b/tests/test_transpilation.py @@ -33,8 +33,9 @@ """ import ast +import sys from textwrap import dedent -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import Mock, MagicMock, sentinel import lark @@ -937,3 +938,29 @@ def test_core_transpile( assert computed == mock_protobuf.return_value else: assert computed == expected_value + + +def test_transpile_function_from_other_module(monkeypatch): + """ + A function from a module the ``evaluation`` module never imports (e.g., ``__main__``) + is resolved when the transpiled code is evaluated. + See https://github.com/cloud-custodian/cel-python/issues/145 + """ + user_module = ModuleType("cel_user_functions") + exec( + dedent("""\ + import celpy.celtypes + + def lower(text): + return celpy.celtypes.StringType(text.lower()) + """), + user_module.__dict__, + ) + monkeypatch.setitem(sys.modules, "cel_user_functions", user_module) + + env = celpy.Environment(runner_class=celpy.CompiledRunner) + expr_ast = env.compile("device.lower() == 'tablet'") + program = env.program(expr_ast, functions={"lower": user_module.lower}) + assert "cel_user_functions.lower(" in program.tp.source_text + result = program.evaluate({"device": celtypes.StringType("TABLET")}) + assert result == celtypes.BoolType(True)