Merge YAML configurations

1
#!/usr/bin/env python
2
"""
3
Backup destination file and merge YAML contents of src into dest.
4
5
By default creates backups, overwrites destination, and clobbers lists.
6
7
Usage:
8
    merge_yaml.py src dest [--create-backup=True] [--dry-run] [--show-stacktrace=False] [--merge-lists=True] [--help]
9
"""
10
11
import argparse
12
import os
13
import shutil
14
from datetime import datetime
15
import sys
16
from pathlib import Path
17
18
# Check Python version
19
if sys.version_info < (3, 0):
20
    print("Error: This script requires Python 3.0 or greater.")
21
    sys.exit(2)
22
23
# Try importing ruamel.yaml
24
try:
25
    from ruamel.yaml import YAML
26
except ImportError:
27
    print(
28
        "Error: ruamel.yaml is not installed. Please install it using 'pip install ruamel.yaml'"
29
    )
30
    sys.exit(2)
31
32
yaml = YAML()
33
34
def deep_merge(d1, d2, **kwargs):
35
    """Deep merges dictionary d2 into dictionary d1."""
36
    merge_lists = kwargs.get("merge_lists")
37
    for key, value in d2.items():
38
        if key in d1:
39
            if isinstance(d1[key], dict) and isinstance(value, dict):
40
                deep_merge(d1[key], value, **kwargs)
41
            elif merge_lists and isinstance(d1[key], list) and isinstance(value, list):
42
                d1[key].extend(value)
43
            else:
44
                d1[key] = value
45
        else:
46
            d1[key] = value
47
    return d1
48
49
def load_yaml_file(filename):
50
    """Load YAML data from a file."""
51
    if not os.path.exists(filename):
52
        return {}
53
    with open(filename, "r") as file:
54
        return yaml.load(file)
55
56
def save_yaml_file(filename, data):
57
    """Save YAML data to a file."""
58
    with open(filename, "w") as file:
59
        yaml.dump(data, file)
60
61
def create_backup(filename):
62
    """Create a timestamped backup of the file."""
63
    # create a directory called backups relative to the filename
64
    backup_dir = filename.parent / "yaml_backups"
65
    try:
66
        backup_dir.mkdir(exist_ok=True)
67
    except Exception as e:
68
        print(
69
            f"Error: Could not create backup directory {backup_dir}. Check your file-permissions or use --no-create-backup to skip creating a backup."
70
        )
71
        exit(2)
72
73
    timestamp = datetime.now().strftime("%y%m%d%H%M%S")
74
    backup_filename = backup_dir / f"{filename.name}.{timestamp}.bak"
75
    shutil.copyfile(filename, backup_filename)
76
    print(f"Backup created: {backup_filename}")
77
78
def main():
79
    parser = argparse.ArgumentParser(
80
        description="Deep merge YAML contents of src into dest."
81
    )
82
    parser.add_argument("src", type=Path, help="Source filename")
83
    parser.add_argument("dest", type=Path, help="Destination filename")
84
    parser.add_argument(
85
        "--create-backup",
86
        type=bool,
87
        default=True,
88
        help="Create a backup of the destination file before merging",
89
    )
90
    parser.add_argument(
91
        "--dry-run",
92
        action="store_true",
93
        help="Print to stdout only, do not write to the destination file",
94
    )
95
    # add a argument for showing the stack trace on yaml parse errors
96
    parser.add_argument(
97
        "--show-stacktrace",
98
        action="store_true",
99
        help="Show stack trace on yaml parse errors",
100
    )
101
    # add an argument to clobber lists
102
    parser.add_argument(
103
        "--merge-lists",
104
        action="store_true",
105
        help="Merge list items instead of clobbering",
106
        default=False,
107
    )
108
109
    args = parser.parse_args()
110
111
    src_filename = args.src.resolve().expanduser()
112
    dest_filename = args.dest.resolve().expanduser()
113
114
    # make sure both files exist
115
    if not src_filename.exists():
116
        print(f"Error: {args.src} does not exist")
117
        exit(2)
118
119
    if not dest_filename.exists():
120
        print(f"Error: {args.dest} does not exist")
121
        exit(2)
122
123
    try:
124
        src_data = load_yaml_file(src_filename)
125
    except Exception as e:
126
        print(
127
            f"Error: {args.src} is not a valid YAML file. Run with --show-stacktrace to see the error."
128
        )
129
        if args.show_stacktrace:
130
            raise e
131
        exit(2)
132
    try:
133
        dest_data = load_yaml_file(dest_filename)
134
    except Exception as e:
135
        print(
136
            f"Error: {args.dest} is not a valid YAML file. Run with --show-stacktrace to see the error."
137
        )
138
        if args.show_stacktrace:
139
            raise e
140
        exit(2)
141
142
    if args.create_backup and not args.dry_run:
143
        create_backup(dest_filename)
144
145
    src_data = load_yaml_file(args.src)
146
    dest_data = load_yaml_file(args.dest)
147
148
    # if dest_data is empty, just copy src_data to dest_data
149
    if not dest_data:
150
        if not args.dry_run:
151
            save_yaml_file(args.dest, src_data)
152
    else:
153
        merged_data = deep_merge(dest_data, src_data, merge_lists=args.merge_lists)
154
        if not args.dry_run:
155
            save_yaml_file(args.dest, merged_data)
156
            print(f"Merged data from {args.src} into {args.dest}")
157
        else:
158
            yaml.dump(merged_data, sys.stdout)
159
160
if __name__ == "__main__":
161
    main()