From fad05632a2c6a05aae0bf3b15366a1d9cccbdc99 Mon Sep 17 00:00:00 2001
From: Alexandre Duret-Lutz <adl@lrde.epita.fr>
Date: Sun, 27 Sep 2015 22:02:49 +0200
Subject: [PATCH] python: implement the map and transform functions for
 formulas

* wrap/python/spot.py: Implement them.
* wrap/python/tests/ltlsimple.py: New tests.
---
 wrap/python/spot.py            | 23 +++++++++++++++++++++++
 wrap/python/tests/ltlsimple.py | 19 +++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/wrap/python/spot.py b/wrap/python/spot.py
index c3b8498c7..d74257890 100644
--- a/wrap/python/spot.py
+++ b/wrap/python/spot.py
@@ -116,9 +116,32 @@ def _formula_to_str(self, format = 'spot', parenth = False):
     else:
         raise ValueError("unknown string format: " + format)
 
+def _formula_traverse(self, func):
+    if func(self):
+        return
+    for f in self:
+        f.traverse(func)
+
+def _formula_map(self, func):
+    k = self.kind()
+    if k in (FalseVal, TrueVal, EmptyWord, AP):
+        return self;
+    if k in (Not, X, F, G, Closure, NegClosure, NegClosureMarked):
+        return formula.unop(k, func(self[0]))
+    if k in (Xor, Implies, Equiv, U, R, W, M, EConcat,
+             EConcatMarked, UConcat):
+        return formula.binop(k, func(self[0]), func(self[1]))
+    if k in (Or, OrRat, And, AndRat, AndNLM, Concat, Fusion):
+        return formula.multop(k, [func(x) for x in self])
+    if k in (Star, FStar):
+        return formula.bunop(k, func(self[0]), self.min(), self.max())
+    raise ValueError("unknown type of formula")
+
 formula.__init__ = _formula_str_ctor
 formula.to_str = _formula_to_str
 formula.show_ast = _render_formula_as_svg
+formula.traverse = _formula_traverse
+formula.map = _formula_map
 
 def _twa_to_str(a, format='hoa', opt=None):
     format = format.lower()
diff --git a/wrap/python/tests/ltlsimple.py b/wrap/python/tests/ltlsimple.py
index 0cf231f50..ba65aa1b4 100755
--- a/wrap/python/tests/ltlsimple.py
+++ b/wrap/python/tests/ltlsimple.py
@@ -74,3 +74,22 @@ assert spot.fnode_instances_check()
 
 #----------------------------------------------------------------------
 assert str([x for x in spot.formula('a &b & c')]) == '[a, b, c]'
+
+
+def switch_g_f(x):
+    if x._is(spot.G):
+        return spot.formula.F(switch_g_f(x[0]))
+    if x._is(spot.F):
+        return spot.formula.G(switch_g_f(x[0]))
+    return x.map(switch_g_f)
+
+f = spot.formula('GFa & XFGb & Fc & G(a | b | Fd)')
+assert str(switch_g_f(f)) == 'FGa & XGFb & Gc & F(a | b | Gd)'
+
+x = 0
+def count_g(f):
+    global x
+    if f._is(spot.G):
+        x += 1
+f.traverse(count_g)
+assert x == 3