python: better way to extend existing classes

* wrap/python/spot.py: Use a decorator to extend classes.
* wrap/python/tests/formulas.ipynb: Adjust expected help text.

wrap/python/spot.py

@@ -23,6 +23,21 @@ import sys
 from functools import lru_cache
 
 
+def _extend(*classes):
+    """
+    Decorator that extends all the given classes with the contents
+    of the class currently being defined.
+    """
+    def wrap(this):
+        for cls in classes:
+            for (name, val) in this.__dict__.items():
+                if name not in ('__dict__', '__weakref__') \
+                   and not (name == '__doc__' and val is None):
+                    setattr(cls, name, val)
+        return classes[0]
+    return wrap
+
+
 def setup(**kwargs):
     """Configure Spot for fancy display.
 
@@ -56,6 +71,7 @@ def setup(**kwargs):
     d = 'rf({})'.format(kwargs.get('font', 'Lato')) + bullets
     os.environ['SPOT_DOTDEFAULT'] = d
 
+
 # In version 3.0.2, Swig puts strongly typed enum in the main
 # namespace without prefixing them.  Latter versions fix this.  So we
 # can remove for following hack once 3.0.2 is no longer used in our
@@ -92,43 +108,71 @@ def _ostream_to_svg(ostr):
     return _str_to_svg(ostr.str().encode('utf-8'))
 
 
-def _render_automaton_as_svg(a, opt=None):
+@_extend(twa, ta)
+class twa:
+    def _repr_svg_(self, opt=None):
+        """Output the automaton as SVG"""
         ostr = ostringstream()
-    print_dot(ostr, a, opt)
+        print_dot(ostr, self, opt)
         return _ostream_to_svg(ostr)
 
-
+    def show(self, opt=None):
-twa._repr_svg_ = _render_automaton_as_svg
+        """Display the automaton as SVG, in the IPython/Jupyter notebook"""
-ta._repr_svg_ = _render_automaton_as_svg
+        # Load the SVG function only if we need it. This way the
+        # bindings can still be used outside of IPython if IPython is
+        # not installed.
+        from IPython.display import SVG
+        return SVG(self._repr_svg_(opt))
 
 
-def _render_formula_as_svg(a):
+@_extend(twa)
+class twa:
+    def to_str(a, format='hoa', opt=None):
+        format = format.lower()
+        if format == 'hoa':
+            ostr = ostringstream()
+            print_hoa(ostr, a, opt)
+            return ostr.str()
+        if format == 'dot':
+            ostr = ostringstream()
+            print_dot(ostr, a, opt)
+            return ostr.str()
+        if format == 'spin':
+            ostr = ostringstream()
+            print_never_claim(ostr, a, opt)
+            return ostr.str()
+        if format == 'lbtt':
+            ostr = ostringstream()
+            print_lbtt(ostr, a, opt)
+            return ostr.str()
+        raise ValueError("unknown string format: " + format)
+
+    def save(a, filename, format='hoa', opt=None, append=False):
+        with open(filename, 'a' if append else 'w') as f:
+            s = a.to_str(format, opt)
+            f.write(s)
+            if s[-1] != '\n':
+                f.write('\n')
+        return a
+
+
+@_extend(formula)
+class formula:
+    def __init__(self, str):
+        """Parse the given string to create a formula."""
+        self.this = parse_formula(str)
+
+    def show_ast(self):
+        """Display the syntax tree of the formula."""
         # Load the SVG function only if we need it. This way the bindings
         # can still be used outside of IPython if IPython is not
         # installed.
         from IPython.display import SVG
         ostr = ostringstream()
-    print_dot_psl(ostr, a)
+        print_dot_psl(ostr, self)
         return SVG(_ostream_to_svg(ostr))
 
-
+    def to_str(self, format='spot', parenth=False):
-def _return_automaton_as_svg(a, opt=None):
-    # Load the SVG function only if we need it. This way the bindings
-    # can still be used outside of IPython if IPython is not
-    # installed.
-    from IPython.display import SVG
-    return SVG(_render_automaton_as_svg(a, opt))
-
-
-twa.show = _return_automaton_as_svg
-ta.show = _return_automaton_as_svg
-
-
-def _formula_str_ctor(self, str):
-    self.this = parse_formula(str)
-
-
-def _formula_to_str(self, format='spot', parenth=False):
         if format == 'spot' or format == 'f':
             return str_psl(self, parenth)
         elif format == 'spin' or format == 's':
@@ -146,8 +190,7 @@ def _formula_to_str(self, format='spot', parenth=False):
         else:
             raise ValueError("unknown string format: " + format)
 
-
+    def __format__(self, spec):
-def _formula_format(self, spec):
         """Format the formula according to `spec`.
 
         Parameters
@@ -221,15 +264,13 @@ def _formula_format(self, spec):
 
         return s.__format__(spec)
 
-
+    def traverse(self, func):
-def _formula_traverse(self, func):
         if func(self):
             return
         for f in self:
             f.traverse(func)
 
-
+    def map(self, func):
-def _formula_map(self, func):
         k = self.kind()
         if k in (op_ff, op_tt, op_eword, op_ap):
             return self
@@ -247,48 +288,6 @@ def _formula_map(self, func):
         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.__format__ = _formula_format
-formula.map = _formula_map
-
-
-def _twa_to_str(a, format='hoa', opt=None):
-    format = format.lower()
-    if format == 'hoa':
-        ostr = ostringstream()
-        print_hoa(ostr, a, opt)
-        return ostr.str()
-    if format == 'dot':
-        ostr = ostringstream()
-        print_dot(ostr, a, opt)
-        return ostr.str()
-    if format == 'spin':
-        ostr = ostringstream()
-        print_never_claim(ostr, a, opt)
-        return ostr.str()
-    if format == 'lbtt':
-        ostr = ostringstream()
-        print_lbtt(ostr, a, opt)
-        return ostr.str()
-    raise ValueError("unknown string format: " + format)
-
-
-def _twa_save(a, filename, format='hoa', opt=None, append=False):
-    with open(filename, 'a' if append else 'w') as f:
-        s = a.to_str(format, opt)
-        f.write(s)
-        if s[-1] != '\n':
-            f.write('\n')
-    return a
-
-
-twa.to_str = _twa_to_str
-twa.save = _twa_save
-
-
 def automata(*filenames):
     """Read automata from a list of sources.
 

wrap/python/tests/formulas.ipynb

@@ -243,9 +243,9 @@
        "output_type": "stream",
        "stream": "stdout",
        "text": [
-        "Help on function _formula_format in module spot:\n",
+        "Help on function __format__ in module spot:\n",
         "\n",
-        "_formula_format(self, spec)\n",
+        "__format__(self, spec)\n",
         "    Format the formula according to `spec`.\n",
         "    \n",
         "    Parameters\n",