generate_state_machine.py

import graphviz
import argparse
from collections import namedtuple
import sys

Transition = namedtuple(
    "Transition", ["src_state", "dst_state", "event", "color"])


def create_graph(transitions, output_path):
    f = graphviz.Digraph('finite_state_machine',
                         node_attr={'color': 'lightblue2', 'style': 'filled', })
    f.attr(rankdir='LR')

    f.attr('node', shape='circle')
    for trans in transitions:
        f.edge(trans.src_state, trans.dst_state,
               trans.event, color=trans.color)
    f.render(output_path, cleanup=True)


def create_h_file(transitions, output_path):
    events = set("EV_"+trans.event for trans in transitions)
    events = sorted(events)
    states = set("S_"+trans.src_state for trans in transitions).union(
        set("S_"+trans.dst_state for trans in transitions))
    states = sorted(states)
    colors = set("ACTION_"+trans.color.upper() for trans in transitions)
    colors = sorted(colors)

    data = ""
    data += f"""/* 
* THIS FILE WAS AUTOGENERATED BY RUNNING: 
* "{' '.join(sys.argv)}".
*/
"""
    data += "#ifndef {0}_H\n#define {0}_H\n".format(output_path.upper())
    data += "#include <stdio.h>\n"
    data += '#include "P_Utils/state_machine.h"\n'
    data += """
typedef enum 
{{
\t{0}
}} sm_state_t;
""".format(",\n\t".join(states))

    data += """
typedef enum
{{
\t{0}
}} sm_event_t;

""".format(",\n\t".join(events))
    data += """
typedef enum
{{
\t{0}
}} sm_action_t;

""".format(",\n\t".join(colors))

    data += "sm_transition_t* get_transition_table(uint16_t* size);\n"
    data += "#endif\n"
    with open(output_path + ".h", 'w') as f:
        f.write(data)

def create_c_file(transitions, output_path):
    data = ""
    data += f"""/* 
* THIS FILE WAS AUTOGENERATED BY RUNNING: 
* "{' '.join(sys.argv)}".
*/
"""
    data += f'#include "{output_path}.h"\n'

    data += "\nstatic sm_transition_t transition_table[] = {\n// clang-format off\n"
    for trans in transitions:
        data += f"\t{{.state=S_{trans.src_state}, .event=EV_{trans.event}, .next_state=S_{trans.dst_state}, .action=ACTION_{trans.color.upper()}}},\n"
    data += "// clang-format on\n};\n"
    data += """
sm_transition_t* get_transition_table(uint16_t* size) {
	*size = sizeof(transition_table)/sizeof(*transition_table);
	return transition_table;
}
"""
    with open(output_path + ".c", 'w') as f:
        f.write(data)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input-file', '-i', type=argparse.FileType('r'),
                        required=True,  help="Path to the serialized automaton. The file should consists of \
                            lines in the format of: \"{SOURCE_STATE} {TARGET_STATE} {EVENT} {COLOR}\"")
    parser.add_argument('--output-image', '-oi',
                        help="The output image filename")
    parser.add_argument('--output-c-files', '-oc',
                        help="The output header and source filename prefix. \
                        For example, providing 'sm' will produce `sm.c` and `sm.h`")
    args = parser.parse_args()

    data = args.input_file.read().replace("-", "_")
    transitions = [Transition(*line.strip().split())
                   for line in data.splitlines() if line and not line.strip().startswith("#")]
    
    if args.output_image:
        create_graph(transitions, args.output_image)
        print(f"Written graph file to '{args.output_image}.pdf'")
    if args.output_c_files:
        create_h_file(transitions, args.output_c_files)
        print(f"Written H file to '{args.output_c_files}.h'")
        create_c_file(transitions, args.output_c_files)
        print(f"Written C file to '{args.output_c_files}.c'")


if __name__ == "__main__":
    main()