Convert an abacus-style computation diagram into a plantUML diagram

focsuml.py

#!/usr/bin/env python3
"""Convert an abacus-style computation diagram into a plantUML diagram
https://regexr.com/3vj09
"""
import argparse
import re
import warnings


parser = argparse.ArgumentParser(description='Generate PlantUML activity diagrams from FoCS abacus-style diagram lists.')
parser.add_argument('infile', type=argparse.FileType('r'),
                    help='input file (use `-` for stdin, the default)')
parser.add_argument('-o', '--outfile', type=argparse.FileType('w'),
                    help='output file (use `-` for stdout')
parser.add_argument('-w', '--web', action='store_true',
                    help='generate diagram on the plantuml.com server (requires `plantuml` package)')
args = parser.parse_args()

DIAGRAM_REGEX = '\(\s*(\d*)\s*,\s*"((?:[^"]|\\.)*)"\s*,\s*"((?:[^"]|\\.)*)"\s*,\s*(-?\d*)\s*,\s*(-?\d*)\s*\)'
DIAGRAM_COMP = re.compile(DIAGRAM_REGEX)

ocaml_text = args.infile.read()

matches = DIAGRAM_COMP.findall(ocaml_text)


def make_uml_path(a, b, label=None):
    template = '{origin} -->'
    if label is not None:
        template += "[{label}]"
    template += ' {dest}'
    return template.format(
        origin=make_uml_entity(a),
        dest=make_uml_entity(b),
        label=label)

def make_uml_entity(name):
    if name in defined:
        return name
    else:
        defined.append(name)
        return '"{}" as {}'.format(nodes[name], name)


def escape_uml(value):
    value = str(value)
    if " " in value or (len(value) > 0 and value[0] == '-'):
        return '"{}"'.format(value)
    else:
        return value


def step_to_uml(this_node, location, action, success_node, failure_node):
    map(escape_uml, [this_node, success_node, failure_node])
    if success_node == "-1":
        success_node = "(*)"
    if failure_node == "-1":
        failure_node = "(*)"
    if this_node == "0":
        this_node = '(*top)'

    if location == "" and action == "":
        return make_uml_path(this_node, success_node)
    elif action == "+":
        return make_uml_path(
            this_node, success_node)
    elif action == "-":
        return [
            make_uml_path(
                this_node, success_node, 'Y'),
            make_uml_path(
                this_node, failure_node, 'N')]


lines = ["@startuml"]
nodes = {escape_uml(match[0]): "**{}**\\n{} {}".format(match[0], match[1], match[2]) for match in matches}  # dict of name: description
defined = ['(*)', '(*top)']  # manages defined
for match in matches:
    try:
        output = step_to_uml(*match)
    except KeyError as e:
        warnings.warn('Unknown node {}'.format(e))
    if isinstance(output, list):
        lines += output
    else:
        lines.append(output)
lines.append("@enduml")

if args.outfile:
    args.outfile.write('\n'.join(lines))

if args.web:
    import plantuml
    server = plantuml.PlantUML()
    print()
    print(server.get_url('\n'.join(lines)))