Convert Docker multi-stage build files into a Mermaid flowchart representation.

dockerfile2mermaid.py

#!/usr/bin/env python3
"""
dockerfile2mermaid

This module provides a utility to convert Docker multi-stage build files into a Mermaid flowchart representation.

Features:
- Extracts all build stages from a Dockerfile.
- Highlights the default build target.
- Represents external images with dashed borders.
- Maps dependencies between stages using different arrow styles based on the Dockerfile directives:
  * `FROM` dependencies are shown with a solid line and a full arrow head.
  * `COPY --from=...` dependencies are shown with a dashed line and an empty arrow head.
  * (Optional) `RUN --mount=(.*)from=...` dependencies, if implemented, can be shown with a dotted line and an empty diamond arrow head.

Usage:
  Command line interface that accepts the Dockerfile as a mandatory argument. By default, it outputs the Mermaid code to stdout.
  An optional `-o` flag allows specifying an output file.

Requirements:
- Python 3.6+
- pathlib

Authors:
  - ChatGPT4 (reviewed by Martín Gaitán)
"""

import argparse
import re
from pathlib import Path


def parse_dockerfile(dockerfile_path: Path):
    stages = {}
    external_images = set()
    stage_counter = 0  # Counter for unnamed stages

    content = dockerfile_path.read_text()

    # Extract all stages with 'AS' directive
    for match in re.finditer(r'(?i)AS\s+(\w+)', content):
        stages[match.group(1)] = []

    # Extract FROM dependencies
    for match in re.finditer(r'(?i)FROM\s+([\w:/.-]+)(?:\s+AS\s+(\w+))?', content):
        if match.group(2):  # If there's an "AS" alias
            stage_name = match.group(2)
        else:
            # Assign a default name to the unnamed stages
            stage_name = f"stage_{stage_counter}"
            stage_counter += 1

        if stage_name not in stages:
            stages[stage_name] = []

        stages[stage_name].append(('solid', match.group(1)))
        if match.group(1) not in stages:
            external_images.add(match.group(1))

    # Extract COPY dependencies
    for match in re.finditer(r'(?i)COPY\s+--from=(\w+|\d+)', content):
        current_stage = list(stages.keys())[-1]

        # Handling the numeric reference
        reference = match.group(1)
        if reference.isdigit():
            reference = list(stages.keys())[int(reference)]

        stages[current_stage].append(("dashed", reference))

    return stages, external_images


def generate_mermaid_code(stages, external_images):
    mermaid_code = "graph TD\n"

    # Define Classes
    mermaid_code += "classDef external stroke-dasharray: 5 5;\n"
    mermaid_code += "classDef default_target fill:#cccccc;\n"

    # External Images
    for image in external_images:
        mermaid_code += f"{image}[{image}]:::external\n"

    # Stages
    default_target = list(stages.keys())[-1]  # Assuming last stage is default
    for stage, dependencies in stages.items():
        style_class = "default_target" if stage == default_target else ""
        for style, dependency in dependencies:
            if style == "solid":
                arrow = "-->"
            elif style == "dashed":
                arrow = "-.->"
            if style_class:
                mermaid_code += f"{dependency} {arrow} {stage}:::{style_class}\n"
            else:
                mermaid_code += f"{dependency} {arrow} {stage}\n"

    return mermaid_code


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate Mermaid flowchart from Dockerfile")
    parser.add_argument("dockerfile", type=Path, help="Path to the Dockerfile")
    parser.add_argument("-o", "--output", type=Path, help="Output file for Mermaid code")

    args = parser.parse_args()

    stages, external_images = parse_dockerfile(args.dockerfile)
    mermaid_code = generate_mermaid_code(stages, external_images)

    if args.output:
        args.output.write_text(mermaid_code)
    else:
        print(mermaid_code)