Passed
Pull Request — master (#12)
by Konstantinos
01:20
created

visualize-dockerfile.generate_markdown()   A

Complexity

Conditions 3

Size

Total Lines 17
Code Lines 11

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 3
eloc 11
nop 2
dl 0
loc 17
rs 9.85
c 0
b 0
f 0
1
#!/usr/bin/env python
2
import argparse
3
import re
4
import typing as t
5
from pathlib import Path
6
7
8
def parse_dockerfile(dockerfile_path):
9
    # Stages built from 'FROM <a> AS <b> statements
10
    stages = {}
11
    # copies built from 'COPY --from=<a> <source_path> <target_path>' statements
12
    copies = {}
13
    # Line match at current line
14
    current_stage = None
15
    stage_name_reg = r'[\w:\.\-]+'
16
    stage_reg = re.compile(
17
        rf'^FROM\s+(?P<stage>{stage_name_reg})\s+[Aa][Ss]\s+(?P<alias>{stage_name_reg})'
18
    )
19
20
    copy_from_reg = re.compile(
21
        r'^COPY\s+\-\-from=(?P<prev_stage>[\w\.\-:]+)\s+(?P<path>[\w\.\-:/]+)'
22
    )
23
24
    with open(dockerfile_path, 'r') as f:
25
        lines = f.readlines()
26
27
        for line in lines:
28
            line = line.strip()
29
30
            # Check if it's a new stage
31
            # each stage has a unique alias in the Dockerfile
32
            stage_match = stage_reg.match(line)
33
            if stage_match:  # FROM <a> AS <b>
34
                current_stage = stage_match.group('alias')
35
36
                # we create an empty list for pointing to "prev" stages
37
                stages[current_stage] = []
38
                copies[current_stage] = []
39
                try:
40
                    previous_stage = stage_match.group('stage')
41
                except AttributeError as error:
42
                    print(f'[DEBUG] Line: {line}')
43
                    print(f"Error: {error}")
44
                    raise error
45
                # Add instructions to current stage
46
                if current_stage:
47
                    stages[current_stage].append(previous_stage)
48
            else:
49
                match = copy_from_reg.match(line)
50
                if match:  # COPY --from=<a> <source_path> <target_path>
51
                    previous_stage: str = match.group('prev_stage')
52
                    path_copied: str = match.group('path')
53
                    copies[current_stage].append((previous_stage, path_copied))
54
55
    return stages, copies
56
57
58
def generate_mermaid_flow_chart(dockerfile_dag):
59
    stages: t.Dict[str, t.List[str]] = dockerfile_dag[0]
60
    thick_line_with_arrow = '-->'
61
    copies: t.Dick[str, t.List[t.Tuple[str, str]]] = dockerfile_dag[1]
62
    dotted_arrow_with_text = '-. "{text}" .->'
63
64
    chart = "graph TB;\n"
65
66
    for stage, prev_stages in stages.items():
67
        # chart += f"  {stage}({stage})\n"
68
69
        # Connect 'FROM <a> AS <b>' Stages
70
        for prev_stage in prev_stages:
71
            chart += f"  {prev_stage} {thick_line_with_arrow} {stage}\n"
72
73
        # Connect 'COPY --from<a> <path> <dest>' statements
74
        prev_copies = copies.get(stage, [])
75
        for prev_copy in prev_copies:
76
            prev_stage: str = prev_copy[0]
77
            # write copied path in arrow text
78
            # path_copied: str = prev_copy[1]
79
            # chart += f"  {prev_stage} " + dotted_arrow_with_text.format(text=path_copied) + f" {stage}\n"
80
            # write COPY (literal) in arrow text
81
            chart += (
82
                f"  {prev_stage} " + dotted_arrow_with_text.format(text='COPY') + f" {stage}\n"
83
            )
84
85
    return chart
86
87
88
def generate_markdown(dockerfile_path, output_path):
89
    dockerfile_dag = parse_dockerfile(dockerfile_path)
90
91
    flow_chart = generate_mermaid_flow_chart(dockerfile_dag)
92
93
    markdown = (
94
        "## Dockerfile Flow Chart\n\n"
95
        f"**Dockerfile: {dockerfile_path}**\n\n"
96
        f"```mermaid\n{flow_chart}```\n"
97
    )
98
    if output_path is None:
99
        print(markdown)
100
        return
101
    with open(output_path, 'w') as f:
102
        f.write(markdown)
103
104
    print(f"Markdown generated and saved to {output_path}")
105
106
107
def parse_cli_args() -> t.Tuple[Path, t.Optional[str]]:
108
    parser = argparse.ArgumentParser(description='Process Dockerfile paths.')
109
110
    parser.add_argument(
111
        'dockerfile_path', nargs='?', default='Dockerfile', help='Path to the Dockerfile'
112
    )
113
    parser.add_argument(
114
        '-o', '--output', help='Output path. If not specified, print to stdout.'
115
    )
116
117
    args = parser.parse_args()
118
119
    dockerfile: Path = Path(args.dockerfile_path)
120
    if not dockerfile.exists():
121
        # explicitly use cwd to try again to find it
122
        dockerfile = Path.cwd() / args.dockerfile_path
123
124
    return dockerfile, args.output
125
126
127
if __name__ == '__main__':
128
    dockerfile_path, output_path = parse_cli_args()
129
    generate_markdown(dockerfile_path, output_path)
130