scripts/bcfg2/pygen.py

#!/usr/bin/env python
# -*- mode: python; coding: utf-8 -*-
#
# pygen.py
# --------
#
# Copyright (C) 2007 Jeremie Dimino <jeremie@dimino.org>
#
# This file is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This file is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Street #330, Boston, MA 02111-1307, USA.

'''Module utilisé par le plugin Python de bcfg2 pour générer un
fichier de configuration à partir d'un script python'''

__all__ = [ "Environment", "compileSource", "generate", "load" ]

import re, marshal, os, sys
from cStringIO import StringIO

# Pour l'arrêt d'un script
class Done(Exception):
    pass

def done():
    raise Done()

class Environment(dict):
    '''Environment dans lequel sont éxécuté les scripts'''

    # Le flux de sortie
    stream = None

    def __init__(self):
        # Création de l'environment initial
        dict.__init__(self, {
                # Les variables à "exporter" dans le fichier produit
                "exports": set([]),
                # Export d'une variable
                "export": self.export,
                # Dernière variable exportée à avoir été définie
                "last_definition": None,
                # Écrit: variable keysep tostring(value)
                "defvar": self.defvar,
                # La convertion en chaîne de charactère
                "tostring": self.tostring,
                # Définition des convertions
                "conv": {bool: {True: "yes", False: "no"},
                         list: lambda l: ", ".join([str(x) for x in l]),
                         tuple: lambda l: ", ".join([str(x) for x in l])},
                # Fonction de base pour imprimer quelque chose
                "out": self.out,
                # Arrêt du script
                "done": done,
                # Le séparateur pour la forme: variable keysep valeur
                "keysep": "=",
                # Le charactère de commentaire
                "comment_start": "#"})

    def __setitem__(self, variable, value):
        '''Lorsqu'on définie une variable, on "l'exporte" dans
        le fichier produit'''
        if variable in self["exports"]:
            self["defvar"](variable, value)
            dict.__setitem__(self, "last_definition", variable)
        dict.__setitem__(self, variable, value)

    def defvar(self, variable, value):
        '''Fonction par défaut pour gérée la définition des variables exportée.
        écrit 'variable = tostring(value)' dans le fichier produit'''
        self["out"]("%s%s%s\n" % (variable, self["keysep"], self["tostring"](value)))

    def out(self, string):
        '''Fonction de base pour écrire une chaine dans le fichier produit'''
        self.stream.write(string)

    def tostring(self, value):
        '''Fonction de convertion objet python -> chaine dans un format sympa'''
        convertor = self["conv"].get(type(value))
        if convertor:
            if type(convertor) == dict:
                return convertor[value]
            else:
                return convertor(value)
        else:
            return str(value)

    # Cette fonction pourrait être écrite directement dans les scripts
    # mais elle est pratique donc on la met pour tout le monde
    def export(self, variable):
        '''Exporte une variable'''
        self["exports"].add(variable)

__re_special_line = re.compile(r"^([ \t]*)(@|%)(.*)$", re.MULTILINE)
__re_affectation = re.compile(r"([a-zA-Z_][a-zA-Z_0-9]*)[ \t]*=")
__re_space_sep = re.compile(r"([^ \t]*)([ \t].*)")

def compileSource(source, filename=""):
    '''Compile un script'''
    # On commence par remplacer les lignes de la forme
    # @xxx par out("xxx")
    newsource = StringIO()
    start = 0
    for m in __re_special_line.finditer(source):
        newsource.write(source[start:m.start()])
        start = m.end()
        newsource.write(m.group(1))
        linetype = m.group(2)
        if linetype == "@":
            line = m.group(3).replace("\\", "\\\\").replace('"', '\\"')
            if line and line[0] == "#":
                newsource.write('out(comment_start + "')
                line = line[1:]
            else:
                newsource.write('out("')
            newsource.write(line)
            newsource.write('\\n")')
        elif linetype == "%":
            line = m.group(3)
            m = __re_affectation.match(line)
            if m:
                varname = m.group(1)
                newsource.write(line)
                newsource.write("; defvar('")
                newsource.write(varname)
                newsource.write("', tostring(")
                newsource.write(varname)
                newsource.write("))\n")
            else:
                m = __re_space_sep.match(line)
                newsource.write("defvar('")
                newsource.write(m.group(1))
                newsource.write("', tostring(")
                newsource.write(m.group(2))
                newsource.write("))\n")
    newsource.write(source[start:])
    return compile(newsource.getvalue(), filename, "exec")

def generate(code, environment=None):
    '''Évalue un script'''
    if type(code) == str:
        code = compileSource(code)
    if not environment:
        environment = Environment()
    environment.stream = StringIO()
    save_stdout = sys.stdout
    sys.stdout = environment.stream
    try:
        exec(code, environment)
    except Done, _:
        pass
    except Exception, exn:
        sys.stdout = save_stdout
        raise exn
    sys.stdout = save_stdout
    return environment.stream.getvalue()

def load(fname, cfname=None):
    '''Charge un script et le compile, en créant/utilisant un fichier de
    cache'''
    if not cfname:
        cfname = fname + 'c'
    if os.path.exists(cfname) and os.stat(fname).st_mtime <= os.stat(cfname).st_mtime:
        code = marshal.load(file(cfname))
    else:
        code = compileSource(file(fname).read(), fname)
        cfile = open(cfname, "w")
        marshal.dump(code, cfile)
        cfile.close()
    return code

if __name__ == "__main__":
    result = generate(load(sys.argv[1]))

    print "resultat:"
    sys.stdout.write(result)