From 9e5433b4c23efd76398caeaa754f66d8829d3806 Mon Sep 17 00:00:00 2001
From: Laurent Perron <lperron@google.com>
Date: Tue, 5 Jan 2016 22:35:02 +0100
Subject: [PATCH] working version of custom python constraint

---
 examples/python/slitherlink.py                | 13 +---
 .../python/constraint_solver.swig             | 70 +++++++++++++++----
 2 files changed, 59 insertions(+), 24 deletions(-)

diff --git a/examples/python/slitherlink.py b/examples/python/slitherlink.py
index 0eb70e94b6..823c564b82 100644
--- a/examples/python/slitherlink.py
+++ b/examples/python/slitherlink.py
@@ -37,14 +37,12 @@ class BooleanSumEven(pywrapcp.PyConstraint):
     self.__vars = vars
     self.__num_possible_true_vars = pywrapcp.RevInteger(0)
     self.__num_always_true_vars = pywrapcp.RevInteger(0)
-    self.__demons = []
 
   def Post(self):
     for i in range(len(self.__vars)):
       v = self.__vars[i]
       if not v.Bound():
-        demon = BooleanSumEvenUpdateDemon(self, i)
-        self.__demons.append(demon)
+        demon = self.Demon(BooleanSumEven.Update, i)
         v.WhenBound(demon)
 
   def InitialPropagate(self):
@@ -100,7 +98,6 @@ class BooleanSumEven(pywrapcp.PyConstraint):
 def SlitherLink(data):
   num_rows = len(data)
   num_columns = len(data[0])
-  to_keep = []
 
   solver = pywrapcp.Solver('slitherlink')
   h_arcs = [[solver.BoolVar('h_arcs[%i][%i]' % (i, j))
@@ -125,15 +122,11 @@ def SlitherLink(data):
 
   for i in range(num_columns):
     column = [h_arcs[j][i] for j in range(num_rows + 1)]
-    ct = BooleanSumEven(solver, column)
-    to_keep.append(ct)
-    solver.Add(ct)
+    solver.Add(BooleanSumEven(solver, column))
 
   for i in range(num_rows):
     row = [v_arcs[j][i] for j in range(num_columns + 1)]
-    ct = BooleanSumEven(solver, row)
-    to_keep.append(ct)
-    solver.Add(ct)
+    solver.Add(BooleanSumEven(solver, row))
 
   all_vars = []
   for row in h_arcs:
diff --git a/src/constraint_solver/python/constraint_solver.swig b/src/constraint_solver/python/constraint_solver.swig
index 12c6f463d1..9e19c4bf16 100644
--- a/src/constraint_solver/python/constraint_solver.swig
+++ b/src/constraint_solver/python/constraint_solver.swig
@@ -174,8 +174,18 @@ static void PyFunctionSolverToVoid(PyObject* pyfunc,
   std::string __str__() {
     return $self->DebugString();
   }
+  %pythoncode {
+  def Add(self, ct):
+    if isinstance(ct, PyConstraint):
+      self.__python_constraints.append(ct)
+    self.AddConstraint(ct)
+  }  // %pythoncode
 }
 
+%feature("pythonappend") operations_research::Solver::Solver %{
+  self.__python_constraints = []
+%}
+
 // Extend IntervalVar to provide a nicer pythonic API for precedence
 // and scheduling constraints. The macros below help do that concisely.
 %define PRECEDENCE_CONSTRAINT(PythonMethodName, CppEnumName)
@@ -708,7 +718,7 @@ namespace operations_research {
 %unignore Solver;
 %unignore Solver::Solver;
 %unignore Solver::~Solver;
-%rename (Add) Solver::AddConstraint;
+%unignore Solver::AddConstraint;
 %unignore Solver::Solve;
 
 // Solver: Decomposed or specialized Solve() API.
@@ -2129,10 +2139,55 @@ class PyDecisionBuilder(DecisionBuilder):
     return "PyDecisionBuilder"
 
 
+class PyDemon(Demon):
+
+  def RunWrapper(self, solver):
+    try:
+      self.Run(solver)
+    except Exception as e:
+      if 'CP Solver fail' in str(e):
+        solver.ShouldFail()
+      else:
+        raise
+
+  def DebugString(self):
+    return "PyDemon"
+
+class PyConstraintDemon(PyDemon):
+
+  def __init__(self, ct, method, delayed, *args):
+    PyDemon.__init__(self)
+    self.__constraint = ct
+    self.__method = method
+    self.__delayed = delayed
+    self.__args = args
+
+  def Run(self, solver):
+    self.__method(self.__constraint, *self.__args)
+
+  def Priority(self):
+    return (pywrapcp.Solver.DELAYED_PRIORITY if self.__delayed
+            else pywrapcp.Solver.NORMAL_PRIORITY)
+
+  def DebugString(self):
+    return 'PyConstraintDemon'
+
+
 class PyConstraint(Constraint):
 
   def __init__(self, solver):
     Constraint.__init__(self, solver)
+    self.__demons = []
+
+  def Demon(self, method, *args):
+    demon = PyConstraintDemon(self, method, False, *args)
+    self.__demons.append(demon)
+    return demon
+
+  def DelayedDemon(self, method, *args):
+    demon = PyConstraintDemon(self, method, True, *args)
+    self.__demons.append(demon)
+    return demon
 
   def InitialPropagateWrapper(self):
     try:
@@ -2147,17 +2202,4 @@ class PyConstraint(Constraint):
     return "PyConstraint"
 
 
-class PyDemon(Demon):
-
-  def RunWrapper(self, solver):
-    try:
-      self.Run(solver)
-    except Exception as e:
-      if 'CP Solver fail' in str(e):
-        solver.ShouldFail()
-      else:
-        raise
-
-  def DebugString(self):
-    return "PyDemon"
 }  // %pythoncode