SQLAlchemy model to simple PlantUML

sql2uml.py

import sys
import argparse
import importlib


def format_field(col):
    if col.foreign_keys or col.primary_key:
        extra = " ({extra})".format(extra="PK" if col.primary_key else "FK")
    else:
        extra = ""
    out = "{} => {}".format(col.name, col.type.__class__.__name__)
    return out + extra


def model_to_uml(module_path, model_name):
    # todo: error handling
    m = importlib.import_module(module_path)
    model = getattr(m, model_name, None)
    if model is None:
        sys.stderr.write("Model not found")
        sys.exit(-1)

    if getattr(model, "__table__", None) is None:
        sys.stderr.write("Not a valid SQLAlchemy model")
        sys.exit(-2)

    fields = "\n".join(
        [format_field(val) for _, val in model.__table__.columns.items()]
    )

    out = """
    @startuml
    skinparam defaultFontSize 18
    skinparam defaultFontName "Source Code Pro"
    hide circle
    map {model} {{
        {fields}
    }}
    @enduml
    """.format(
        model=model_name, fields=fields
    )
    sys.stdout.write(out)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "model", type=str, help="SQLAlchemy model name to dump as PlantUML entity"
    )
    parser.add_argument(
        "path", type=str, help="Path to python module containing the SQLAlchemy model"
    )
    args = parser.parse_args()
    model_to_uml(args.path, args.model)


if __name__ == "__main__":
    main()