From 19251e5e61f30e0a63b5ba110a19332158d4d363 Mon Sep 17 00:00:00 2001
From: "John J. Camilleri" <john@digitalgrammars.com>
Date: Mon, 27 Sep 2021 14:22:13 +0200
Subject: [PATCH] Add exprProbability

---
 src/runtime/python/pypgf.c                    | 19 +++++++++++++++++++
 src/runtime/python/tests/test_transactions.py | 13 +++++++++----
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/runtime/python/pypgf.c b/src/runtime/python/pypgf.c
index e6bc98ba9..939ad44ff 100644
--- a/src/runtime/python/pypgf.c
+++ b/src/runtime/python/pypgf.c
@@ -318,6 +318,22 @@ PGF_functionProbability(PGFObject *self, PyObject *args)
     return PyFloat_FromDouble((double)prob);
 }
 
+static PyObject *
+PGF_exprProbability(PGFObject *self, PyObject *args)
+{
+    ExprObject *expr;
+    if (!PyArg_ParseTuple(args, "O!", &pgf_ExprType, &expr))
+        return NULL;
+
+    PgfExn err;
+    prob_t prob = pgf_expr_prob(self->db, self->revision, (PgfExpr) expr, &marshaller, &err);
+    if (handleError(err) != PGF_EXN_NONE) {
+        return NULL;
+    }
+
+    return PyFloat_FromDouble((double)prob);
+}
+
 static PyGetSetDef PGF_getseters[] = {
     {"abstractName",
      (getter)PGF_getAbstractName, NULL,
@@ -364,6 +380,9 @@ static PyMethodDef PGF_methods[] = {
     {"functionProbability", (PyCFunction)PGF_functionProbability, METH_VARARGS,
      "Returns the probability of a function"
     },
+    {"exprProbability", (PyCFunction)PGF_exprProbability, METH_VARARGS,
+     "Returns the probability of an expression"
+    },
 
     {"checkoutBranch", (PyCFunction)PGF_checkoutBranch, METH_VARARGS,
      "Switch to a branch"
diff --git a/src/runtime/python/tests/test_transactions.py b/src/runtime/python/tests/test_transactions.py
index 29203c7f1..00cfaab00 100644
--- a/src/runtime/python/tests/test_transactions.py
+++ b/src/runtime/python/tests/test_transactions.py
@@ -76,6 +76,11 @@ def test_original_category_prob(gr1):
     #     gr1.categoryProbability("Q")
     assert gr1.categoryProbability("Q") == float('inf')
 
+def test_original_expr_prob(gr1):
+    # with pytest.raises(KeyError):
+    #     gr1.functionProbability("foo")
+    assert gr1.exprProbability(ExprFun("foo")) == float('inf')
+
 # gr2
 
 def test_extended_functions(gr2):
@@ -91,13 +96,13 @@ def test_extended_function_type(gr2):
     assert gr2.functionType("foo") == ty
 
 def test_extended_function_prob(gr2):
-    # TODO: can't we get higher precision?
-    # assert gr2.functionProbability("foo") == prob
     assert math.isclose(gr2.functionProbability("foo"), prob, rel_tol=1e-06)
 
-@pytest.mark.skip(reason="failing")
 def test_extended_category_prob(gr2):
-    assert gr2.categoryProbability("Q") == prob
+    assert math.isclose(gr2.categoryProbability("Q"), prob, rel_tol=1e-06)
+
+def test_extended_expr_prob(gr2):
+    assert math.isclose(gr2.exprProbability(ExprFun("foo")), prob, rel_tol=1e-06)
 
 # gr3