From c1d09c7a227fe974f653ab820ee547f0d3481745 Mon Sep 17 00:00:00 2001 From: Paulo Lieuthier Date: Fri, 13 Feb 2026 18:21:59 -0300 Subject: [PATCH 01/10] feat: implement add_extension in Python binding --- bindings/python/src/lib.rs | 18 ++++++++++++++++++ bindings/python/test.py | 22 ++++++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 237998c5..16eaee21 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -344,6 +344,24 @@ impl Engine { v.to_json_str() } + pub fn add_extension(&mut self, path: String, nargs: u8, extension: Py) -> Result<()> { + let func_ref = Arc::new(extension); + + let extension_impl = move |args: Vec| -> Result { + Python::with_gil(|py| { + let py_args_vec: Result> = + args.into_iter().map(|arg| to(arg, py)).collect(); + let py_args = PyTuple::new(py, py_args_vec?)?; + let py_result = func_ref.call1(py, py_args)?; + let rego_result = from(&py_result.into_bound(py))?; + Ok(rego_result) + }) + }; + + self.engine + .add_extension(path, nargs, Box::new(extension_impl)) + } + /// Enable code coverage /// /// * `enable`: Whether to enable coverage or not. diff --git a/bindings/python/test.py b/bindings/python/test.py index a1e20259..d7f84512 100644 --- a/bindings/python/test.py +++ b/bindings/python/test.py @@ -163,3 +163,25 @@ def run_host_await_example(): print(vm.resume('{"tier":"gold"}')) run_host_await_example() + +def custom_function(arg1, arg2): + return f"{arg1}, {arg2}!" + +# Extension example +policy = """ +package demo + +result := greeting(a, b) if { + a := data.a + b := data.b +} +""" +def run_extension_example(): + rego = regorus.Engine() + rego.add_policy("demo", policy) + rego.add_extension("greeting", 2, custom_function) + + rego.add_data({"a": "Hello", "b": "World"}) + print(rego.eval_rule("data.demo.result")) + +run_extension_example() From c79564f5a6c4d3848c74e61e66b7410a4beac20f Mon Sep 17 00:00:00 2001 From: Paulo Lieuthier Date: Wed, 18 Feb 2026 20:46:01 -0300 Subject: [PATCH 02/10] doc: add documentation to Python's add_extension --- bindings/python/src/lib.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 16eaee21..95dc8b65 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -344,6 +344,18 @@ impl Engine { v.to_json_str() } + /// Registers a custom Python function as a Rego extension. + /// + /// This allows you to define functions in Python that can be called directly + /// from your Rego policies. The Python function will be called synchronously + /// during policy evaluation. + /// + /// Arguments passed from Rego are automatically converted to their corresponding + /// Python types. The return value is converted back to a Rego value. + /// + /// * `path`: Full path to the function as it will be used in Rego. + /// * `nargs`: The number of arguments the function expects. + /// * `extension`: The Python function to execute. Must accept exactly `nargs` arguments. pub fn add_extension(&mut self, path: String, nargs: u8, extension: Py) -> Result<()> { let func_ref = Arc::new(extension); From c40e910f5de99efd58a3801c32ac6f4053620037 Mon Sep 17 00:00:00 2001 From: Paulo Lieuthier Date: Wed, 18 Feb 2026 20:46:29 -0300 Subject: [PATCH 03/10] fix: check if extension is callable in Python's add_extension --- bindings/python/src/lib.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 95dc8b65..bdd05501 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -361,6 +361,9 @@ impl Engine { let extension_impl = move |args: Vec| -> Result { Python::with_gil(|py| { + if !func_ref.bind(py).is_callable() { + return Err(anyhow!("extension must be callable")) + } let py_args_vec: Result> = args.into_iter().map(|arg| to(arg, py)).collect(); let py_args = PyTuple::new(py, py_args_vec?)?; From a89e38f8693e6806ea0053c0056098546103c0d1 Mon Sep 17 00:00:00 2001 From: Paulo Lieuthier Date: Wed, 18 Feb 2026 20:47:14 -0300 Subject: [PATCH 04/10] feat: provide more information about exceptions from Python extensions --- bindings/python/src/lib.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index bdd05501..e13159a8 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -358,6 +358,7 @@ impl Engine { /// * `extension`: The Python function to execute. Must accept exactly `nargs` arguments. pub fn add_extension(&mut self, path: String, nargs: u8, extension: Py) -> Result<()> { let func_ref = Arc::new(extension); + let path_clone = path.clone(); let extension_impl = move |args: Vec| -> Result { Python::with_gil(|py| { @@ -367,7 +368,8 @@ impl Engine { let py_args_vec: Result> = args.into_iter().map(|arg| to(arg, py)).collect(); let py_args = PyTuple::new(py, py_args_vec?)?; - let py_result = func_ref.call1(py, py_args)?; + let py_result = func_ref.call1(py, py_args) + .map_err(|e| anyhow!("extension '{}' raises Python error: {}", path_clone, e))?; let rego_result = from(&py_result.into_bound(py))?; Ok(rego_result) }) From df142ac42ebe68d411aea04a7d629215d54f3d71 Mon Sep 17 00:00:00 2001 From: Paulo Lieuthier Date: Wed, 18 Feb 2026 20:47:45 -0300 Subject: [PATCH 05/10] test: add more test cases for Python's add_extension --- bindings/python/test.py | 213 +++++++++++++++++++++++++++++++++++++--- 1 file changed, 199 insertions(+), 14 deletions(-) diff --git a/bindings/python/test.py b/bindings/python/test.py index d7f84512..b63563a0 100644 --- a/bindings/python/test.py +++ b/bindings/python/test.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +import typing +import pytest import regorus import sys @@ -164,24 +166,207 @@ def run_host_await_example(): run_host_await_example() -def custom_function(arg1, arg2): - return f"{arg1}, {arg2}!" +def test_extension_execution(): + rego = regorus.Engine() + rego.add_policy("demo", + """ + package demo + + result := greeting(a, b) if { + a := data.a + b := data.b + } + """) + + def custom_function(arg1, arg2): + return f"{arg1}, {arg2}!" + rego.add_extension("greeting", 2, custom_function) -# Extension example -policy = """ -package demo + rego.add_data({"a": "Hello", "b": "World"}) + result = rego.eval_rule("data.demo.result") + assert result == "Hello, World!", f"Unexpected result: {result}" -result := greeting(a, b) if { - a := data.a - b := data.b -} -""" -def run_extension_example(): +def test_extension_wrong_arity(): + rego = regorus.Engine() + rego.add_policy("demo", + """ + package demo + + result := greeting(a, b) if { + a := data.a + b := data.b + } + """) + + def custom_function(arg1, arg2): + return f"{arg1}, {arg2}!" + + rego.add_extension("greeting", 3, custom_function) + rego.add_data({"a": "Hello", "b": "World"}) + + with pytest.raises(RuntimeError) as ex: + rego.eval_rule("data.demo.result") + + assert "error: incorrect number of parameters supplied to extension" in str(ex.value) + +def test_extension_raises_exception(): rego = regorus.Engine() - rego.add_policy("demo", policy) + rego.add_policy("demo", + """ + package demo + + result := greeting(a, b) if { + a := data.a + b := data.b + } + """) + + def custom_function(arg1, arg2): + raise RuntimeError("unknown error") + rego.add_extension("greeting", 2, custom_function) + rego.add_data({"a": "Hello", "b": "World"}) + + with pytest.raises(RuntimeError) as ex: + rego.eval_rule("data.demo.result") + + assert "error: extension 'greeting' raise Python error: RuntimeError: unknown error" in str(ex.value) + +def test_extension_zero_arg(): + rego = regorus.Engine() + rego.add_policy("demo", + """ + package demo + + result := greeting() + """) + def custom_function(): + return "Hello, World!" + + rego.add_extension("greeting", 0, custom_function) + rego.add_data({"a": "Hello", "b": "World"}) + + result = rego.eval_rule("data.demo.result") + assert result == "Hello, World!", f"Unexpected result: {result}" + +def test_extension_non_callable(): + rego = regorus.Engine() + rego.add_policy("demo", + """ + package demo + + result := greeting() + """) + + rego.add_extension("greeting", 0, 123) rego.add_data({"a": "Hello", "b": "World"}) - print(rego.eval_rule("data.demo.result")) -run_extension_example() + with pytest.raises(RuntimeError) as ex: + rego.eval_rule("data.demo.result") + + assert "error: extension must be callable" in str(ex.value) + + +def test_extension_duplicate(): + rego = regorus.Engine() + rego.add_policy("demo", + """ + package demo + + result := greeting() + """) + + def custom_function1(arg1, arg2): + return f"{arg1}, {arg2}!" + def custom_function2(arg1, arg2): + return f"{arg1}, {arg2}!" + + rego.add_extension("greeting", 0, custom_function1) + + with pytest.raises(RuntimeError) as ex: + rego.add_extension("greeting", 0, custom_function2) + + assert "extension already added" in str(ex.value) + + +def test_extension_types(): + rego = regorus.Engine() + rego.add_policy("demo", + """ + package demo + + i := custom.triple(10) + f := custom.triple(2.5) + b1 := custom.negate(true) + b2 := custom.negate(false) + + a := custom.first([true, null, 1]) + b := custom.first([null, null, 1]) + c := custom.first([null, null, null]) + + object := custom.modify_object({"a": 1, "b": 2}) + list := custom.modify_list([3, 4]) + set := custom.modify_set({5, 6}) + """) + + def triple(n): + return n*3 + + def negate(b): + return not b + + def first(list): + for i in list: + if i is not None: + return i + return None + + def modify_object(object): + assert isinstance(object, dict) + return {k: v*2 for k, v in object.items()} + + def modify_list(list): + assert isinstance(list, typing.List) + return [x*2 for x in list] + + def modify_set(set): + assert isinstance(set, typing.Set) + return {x*2 for x in set} + + rego.add_extension("custom.triple", 1, triple) + rego.add_extension("custom.negate", 1, negate) + rego.add_extension("custom.first", 1, first) + rego.add_extension("custom.modify_object", 1, modify_object) + rego.add_extension("custom.modify_list", 1, modify_list) + rego.add_extension("custom.modify_set", 1, modify_set) + + i = rego.eval_rule("data.demo.i") + assert i == 30, f"Unexpected result for 'i': {i}" + + f = rego.eval_rule("data.demo.f") + assert f == 7.5, f"Unexpected result for 'f': {f}" + + b1 = rego.eval_rule("data.demo.b1") + assert b1 == False, f"Unexpected result for 'b1': {b1}" + + b2 = rego.eval_rule("data.demo.b2") + assert b2 == True, f"Unexpected result for 'b2': {b2}" + + a = rego.eval_rule("data.demo.a") + assert a == True, f"Unexpected result for 'a': {a}" + + b = rego.eval_rule("data.demo.b") + assert b == 1, f"Unexpected result for 'b': {b}" + + c = rego.eval_rule("data.demo.c") + assert c is None, f"Unexpected result for 'c': {c}" + + obj = rego.eval_rule("data.demo.object") + assert obj == {"a": 2, "b": 4}, f"Unexpected object: {obj}" + + list = rego.eval_rule("data.demo.list") + assert list == [6, 8], f"Unexpected list: {list}" + + set = rego.eval_rule("data.demo.set") + assert set == {10, 12}, f"Unexpected list: {set}" From 36eda225ee655e1aacd274b9eafd99485d161218 Mon Sep 17 00:00:00 2001 From: Paulo Lieuthier Date: Tue, 24 Feb 2026 09:58:19 -0300 Subject: [PATCH 06/10] feat(python): don't check if extension is callable in every invocation --- bindings/python/src/lib.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index e13159a8..e5094545 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -357,14 +357,18 @@ impl Engine { /// * `nargs`: The number of arguments the function expects. /// * `extension`: The Python function to execute. Must accept exactly `nargs` arguments. pub fn add_extension(&mut self, path: String, nargs: u8, extension: Py) -> Result<()> { + Python::with_gil(|py| { + if !extension.bind(py).is_callable() { + return Err(anyhow!("extension '{}' must be callable", path)); + } + Ok(()) + })?; + let func_ref = Arc::new(extension); let path_clone = path.clone(); let extension_impl = move |args: Vec| -> Result { Python::with_gil(|py| { - if !func_ref.bind(py).is_callable() { - return Err(anyhow!("extension must be callable")) - } let py_args_vec: Result> = args.into_iter().map(|arg| to(arg, py)).collect(); let py_args = PyTuple::new(py, py_args_vec?)?; From 42bcc0be05d76eb6b8d00cbb13ed0e399a1522ee Mon Sep 17 00:00:00 2001 From: Paulo Lieuthier Date: Tue, 24 Feb 2026 10:46:31 -0300 Subject: [PATCH 07/10] test(python): ditch pytest, use manual try-except to assert exceptions, call test methods --- bindings/python/test.py | 71 +++++++++++++++++++++++++---------------- 1 file changed, 44 insertions(+), 27 deletions(-) diff --git a/bindings/python/test.py b/bindings/python/test.py index b63563a0..7f1921b1 100644 --- a/bindings/python/test.py +++ b/bindings/python/test.py @@ -1,8 +1,5 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -import typing -import pytest - import regorus import sys @@ -186,6 +183,8 @@ def custom_function(arg1, arg2): result = rego.eval_rule("data.demo.result") assert result == "Hello, World!", f"Unexpected result: {result}" +test_extension_execution() + def test_extension_wrong_arity(): rego = regorus.Engine() rego.add_policy("demo", @@ -204,10 +203,14 @@ def custom_function(arg1, arg2): rego.add_extension("greeting", 3, custom_function) rego.add_data({"a": "Hello", "b": "World"}) - with pytest.raises(RuntimeError) as ex: + try: rego.eval_rule("data.demo.result") + except RuntimeError as ex: + assert "error: incorrect number of parameters supplied to extension" in str(str(ex)) + else: + assert False, "exception not thrown" - assert "error: incorrect number of parameters supplied to extension" in str(ex.value) +test_extension_wrong_arity() def test_extension_raises_exception(): rego = regorus.Engine() @@ -227,10 +230,15 @@ def custom_function(arg1, arg2): rego.add_extension("greeting", 2, custom_function) rego.add_data({"a": "Hello", "b": "World"}) - with pytest.raises(RuntimeError) as ex: + try: rego.eval_rule("data.demo.result") + except RuntimeError as ex: + assert "error: extension 'greeting' raises Python error: RuntimeError: unknown error" in str(ex) + else: + assert False, "exception not thrown" + +test_extension_raises_exception() - assert "error: extension 'greeting' raise Python error: RuntimeError: unknown error" in str(ex.value) def test_extension_zero_arg(): rego = regorus.Engine() @@ -250,6 +258,8 @@ def custom_function(): result = rego.eval_rule("data.demo.result") assert result == "Hello, World!", f"Unexpected result: {result}" +test_extension_zero_arg() + def test_extension_non_callable(): rego = regorus.Engine() rego.add_policy("demo", @@ -259,13 +269,14 @@ def test_extension_non_callable(): result := greeting() """) - rego.add_extension("greeting", 0, 123) - rego.add_data({"a": "Hello", "b": "World"}) - - with pytest.raises(RuntimeError) as ex: - rego.eval_rule("data.demo.result") + try: + rego.add_extension("greeting", 0, 123) + except RuntimeError as ex: + assert "extension 'greeting' must be callable" in str(ex) + else: + assert False, "exception not thrown" - assert "error: extension must be callable" in str(ex.value) +test_extension_non_callable() def test_extension_duplicate(): @@ -284,10 +295,14 @@ def custom_function2(arg1, arg2): rego.add_extension("greeting", 0, custom_function1) - with pytest.raises(RuntimeError) as ex: + try: rego.add_extension("greeting", 0, custom_function2) + except RuntimeError as ex: + assert "extension already added" in str(ex) + else: + assert False, "exception not thrown" - assert "extension already added" in str(ex.value) +test_extension_duplicate() def test_extension_types(): @@ -316,8 +331,8 @@ def triple(n): def negate(b): return not b - def first(list): - for i in list: + def first(lst): + for i in lst: if i is not None: return i return None @@ -326,13 +341,13 @@ def modify_object(object): assert isinstance(object, dict) return {k: v*2 for k, v in object.items()} - def modify_list(list): - assert isinstance(list, typing.List) - return [x*2 for x in list] + def modify_list(lst): + assert isinstance(lst, list) + return [x*2 for x in lst] - def modify_set(set): - assert isinstance(set, typing.Set) - return {x*2 for x in set} + def modify_set(st): + assert isinstance(st, set) + return {x*2 for x in st} rego.add_extension("custom.triple", 1, triple) rego.add_extension("custom.negate", 1, negate) @@ -365,8 +380,10 @@ def modify_set(set): obj = rego.eval_rule("data.demo.object") assert obj == {"a": 2, "b": 4}, f"Unexpected object: {obj}" - list = rego.eval_rule("data.demo.list") - assert list == [6, 8], f"Unexpected list: {list}" + lst = rego.eval_rule("data.demo.list") + assert lst == [6, 8], f"Unexpected list: {lst}" + + st = rego.eval_rule("data.demo.set") + assert st == {10, 12}, f"Unexpected lst: {st}" - set = rego.eval_rule("data.demo.set") - assert set == {10, 12}, f"Unexpected list: {set}" +test_extension_types() From c88c93f04b8f62dc04897bfb3b0e621c6405afc4 Mon Sep 17 00:00:00 2001 From: Paulo Lieuthier Date: Wed, 25 Feb 2026 00:22:57 -0300 Subject: [PATCH 08/10] doc(python): explain add_extension's clone semantics --- bindings/python/src/lib.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index e5094545..6eff10ad 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -356,6 +356,9 @@ impl Engine { /// * `path`: Full path to the function as it will be used in Rego. /// * `nargs`: The number of arguments the function expects. /// * `extension`: The Python function to execute. Must accept exactly `nargs` arguments. + /// + /// Note: When the engine is cloned, extensions share the same Python callable reference + /// rather than being deep-copied. Stateful callables will share state across clones. pub fn add_extension(&mut self, path: String, nargs: u8, extension: Py) -> Result<()> { Python::with_gil(|py| { if !extension.bind(py).is_callable() { From bdef543fdf800094f4018923bf6636d1f636049a Mon Sep 17 00:00:00 2001 From: Paulo Lieuthier Date: Wed, 25 Feb 2026 00:24:27 -0300 Subject: [PATCH 09/10] python: fix a couple of nits in tests --- bindings/python/test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bindings/python/test.py b/bindings/python/test.py index 7f1921b1..f1b56a72 100644 --- a/bindings/python/test.py +++ b/bindings/python/test.py @@ -206,7 +206,7 @@ def custom_function(arg1, arg2): try: rego.eval_rule("data.demo.result") except RuntimeError as ex: - assert "error: incorrect number of parameters supplied to extension" in str(str(ex)) + assert "error: incorrect number of parameters supplied to extension" in str(ex) else: assert False, "exception not thrown" @@ -384,6 +384,6 @@ def modify_set(st): assert lst == [6, 8], f"Unexpected list: {lst}" st = rego.eval_rule("data.demo.set") - assert st == {10, 12}, f"Unexpected lst: {st}" + assert st == {10, 12}, f"Unexpected set: {st}" test_extension_types() From c9625b277c955f6d1f5985eeba942187cb9b6f58 Mon Sep 17 00:00:00 2001 From: Paulo Lieuthier Date: Wed, 25 Feb 2026 09:26:10 -0300 Subject: [PATCH 10/10] python: format code --- bindings/python/src/lib.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 6eff10ad..858c9026 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -357,7 +357,7 @@ impl Engine { /// * `nargs`: The number of arguments the function expects. /// * `extension`: The Python function to execute. Must accept exactly `nargs` arguments. /// - /// Note: When the engine is cloned, extensions share the same Python callable reference + /// Note: When the engine is cloned, extensions share the same Python callable reference /// rather than being deep-copied. Stateful callables will share state across clones. pub fn add_extension(&mut self, path: String, nargs: u8, extension: Py) -> Result<()> { Python::with_gil(|py| { @@ -375,8 +375,9 @@ impl Engine { let py_args_vec: Result> = args.into_iter().map(|arg| to(arg, py)).collect(); let py_args = PyTuple::new(py, py_args_vec?)?; - let py_result = func_ref.call1(py, py_args) - .map_err(|e| anyhow!("extension '{}' raises Python error: {}", path_clone, e))?; + let py_result = func_ref.call1(py, py_args).map_err(|e| { + anyhow!("extension '{}' raises Python error: {}", path_clone, e) + })?; let rego_result = from(&py_result.into_bound(py))?; Ok(rego_result) })