Signed-off-by: Ján Tomko <jtomko@redhat.com>
---
 scripts/group-qemu-caps.py | 50 ++++++++++++++++++++++----------------
 1 file changed, 29 insertions(+), 21 deletions(-)

diff --git a/scripts/group-qemu-caps.py b/scripts/group-qemu-caps.py
index 34823931df..e4161f5bb5 100755
--- a/scripts/group-qemu-caps.py
+++ b/scripts/group-qemu-caps.py
@@ -61,14 +61,9 @@ def load_caps_flags(lines):
     return capsflags
 
 
-def regroup_caps(check, filename, start_regex, end_regex,
+def regroup_caps(original, start_regex, end_regex,
                  trailing_newline, counter_prefix, capsflags):
     step = 5
-
-    original = []
-    with open(filename, "r") as fh:
-        original = fh.read().splitlines()
-
     fixed = []
     game_on = False
     counter = 0
@@ -107,7 +102,10 @@ def regroup_caps(check, filename, start_regex, end_regex,
                 line = flagname[0] + " /* %s */" % capsflags[counter - 1]
 
         fixed.append(line)
+    return fixed
 
+
+def check(check, filename, original, fixed):
     if check:
         orig = "\n".join(original) + "\n"
         new = "\n".join(fixed) + "\n"
@@ -143,22 +141,32 @@ errs = False
                                        'QEMU_CAPS_LAST /* this must')
 capsflags = load_caps_flags(header_lines[start:end])
 
-if not regroup_caps(args.check,
-                    args.prefix + 'src/qemu/qemu_capabilities.c',
-                    r'virQEMUCaps grouping marker',
-                    r'\);',
-                    0,
-                    "              ",
-                    capsflags):
+(source_lines, source_start, source_end) = load_file(args.prefix + 'src/qemu/qemu_capabilities.c',
+                                                     r'virQEMUCaps grouping marker',
+                                                     r'\);')
+fixed_source = regroup_caps(source_lines,
+                            r'virQEMUCaps grouping marker',
+                            r'\);',
+                            0,
+                            "              ",
+                            capsflags)
+
+fixed_header = regroup_caps(header_lines,
+                            r'virQEMUCapsFlags grouping marker',
+                            r'QEMU_CAPS_LAST \/\* this must',
+                            1,
+                            "    ",
+                            None)
+
+if not check(args.check,
+             args.prefix + 'src/qemu/qemu_capabilities.c',
+             source_lines,
+             fixed_source):
     errs = True
-
-if not regroup_caps(args.check,
-                    args.prefix + 'src/qemu/qemu_capabilities.h',
-                    r'virQEMUCapsFlags grouping marker',
-                    r'QEMU_CAPS_LAST \/\* this must',
-                    1,
-                    "    ",
-                    None):
+if not check(args.check,
+             args.prefix + 'src/qemu/qemu_capabilities.h',
+             header_lines,
+             fixed_header):
     errs = True
 
 if errs:
-- 
2.45.2