Resolve user function modules in CompiledRunner (#145)

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):
     """

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)