From d79e5f4df3b9e85769ecc8b819dc59d7e07bc24c Mon Sep 17 00:00:00 2001
From: wout4 <wout.vanroose@gmail.com>
Date: Fri, 17 May 2024 15:08:05 +0200
Subject: [PATCH] add is_leaf and has_nested to expressions

---
 cpmpy/expressions/core.py      | 11 ++++++++++-
 cpmpy/expressions/utils.py     |  4 ++++
 cpmpy/expressions/variables.py |  8 +++++++-
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/cpmpy/expressions/core.py b/cpmpy/expressions/core.py
index 1b210e680..9c1a88f3d 100644
--- a/cpmpy/expressions/core.py
+++ b/cpmpy/expressions/core.py
@@ -72,7 +72,7 @@
 import numpy as np
 
 
-from .utils import is_num, is_any_list, flatlist, argval, get_bounds, is_boolexpr, is_true_cst, is_false_cst
+from .utils import is_num, is_any_list, flatlist, argval, get_bounds, is_boolexpr, is_true_cst, is_false_cst, is_leaf
 from ..exceptions import IncompleteFunctionError, TypeError
 
 
@@ -135,6 +135,12 @@ def __repr__(self):
     def __hash__(self):
         return hash(self.__repr__())
 
+    def has_nested(self):
+        return not all([is_leaf(x) for x in self.args])
+
+    def is_leaf(self):
+        return False  # default
+
     def is_bool(self):
         """ is it a Boolean (return type) Operator?
             Default: yes
@@ -381,6 +387,9 @@ def __bool__(self):
         """Called to implement truth value testing and the built-in operation bool(), return stored value"""
         return self.args[0]
 
+    def is_leaf(self):
+        return True
+
 
 class Comparison(Expression):
     """Represents a comparison between two sub-expressions
diff --git a/cpmpy/expressions/utils.py b/cpmpy/expressions/utils.py
index a42e88c82..75c4f4df3 100644
--- a/cpmpy/expressions/utils.py
+++ b/cpmpy/expressions/utils.py
@@ -128,6 +128,10 @@ def argval(a):
         raise e
 
 
+def is_leaf(a):
+    return a.is_leaf() if hasattr(a, 'is_leaf') else True
+
+
 def eval_comparison(str_op, lhs, rhs):
     """
         Internal function: evaluates the textual `str_op` comparison operator
diff --git a/cpmpy/expressions/variables.py b/cpmpy/expressions/variables.py
index 31c1ab33b..c61d9278b 100644
--- a/cpmpy/expressions/variables.py
+++ b/cpmpy/expressions/variables.py
@@ -53,7 +53,7 @@
 
 import numpy as np
 from .core import Expression, Operator
-from .utils import is_num, is_int, flatlist, is_boolexpr, is_true_cst, is_false_cst, get_bounds
+from .utils import is_num, is_int, flatlist, is_boolexpr, is_true_cst, is_false_cst, get_bounds, is_leaf
 
 
 def BoolVar(shape=1, name=None):
@@ -245,6 +245,9 @@ def is_bool(self):
         """
         return False
 
+    def is_leaf(self):
+        return True
+
     def value(self):
         """ the value obtained in the last solve call
             (or 'None')
@@ -385,6 +388,9 @@ def is_bool(self):
         """
         return False
 
+    def is_leaf(self):
+        return all([is_leaf(x) for x in self])
+
     def value(self):
         """ the values, for each of the stored variables, obtained in the last solve call
             (or 'None')