From fa67361a68b5c2c90cb4f84a8341dab2d276c5b9 Mon Sep 17 00:00:00 2001
From: Mathieu Comandon <strycore@gmail.com>
Date: Tue, 25 Oct 2016 14:38:29 -0700
Subject: [PATCH] Add render method to registry

---
 lutris/util/wineregistry.py | 29 +++++++++++++++++++++++++----
 tests/test_registry.py      | 20 ++++++++++++++++++--
 2 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/lutris/util/wineregistry.py b/lutris/util/wineregistry.py
index d4bce6e02..0d72e3fc3 100644
--- a/lutris/util/wineregistry.py
+++ b/lutris/util/wineregistry.py
@@ -50,9 +50,13 @@ class WindowsFileTime:
 
 
 class WineRegistry(object):
+    version_header = "WINE REGISTRY Version "
+    relative_to_header = ";; All keys relative to "
+
     def __init__(self, reg_filename=None):
         self.arch = 'win32'
         self.version = 2
+        self.relative_to = "\\\\User\\\\S-1-5-21-0-0-0-1000"
         self.keys = OrderedDict()
         if reg_filename:
             self.prefix_path = os.path.dirname(reg_filename)
@@ -68,21 +72,38 @@ class WineRegistry(object):
         registry_lines = self.get_raw_registry(reg_filename)
         current_key = None
         for line in registry_lines:
-            # Remove trailing newlines
-            line = line.rstrip('\n')
-            if line.startswith('WINE REGISTRY Version'):
-                self.version = int(line[len('WINE REGISTRY Version '):])
+            line = line.rstrip('\n')  # Remove trailing newlines
+
+            if line.startswith(self.version_header):
+                self.version = int(line[len(self.version_header):])
                 continue
+
+            if line.startswith(self.relative_to_header):
+                self.relative_to = line[len(self.relative_to_header):]
+                continue
+
             if line.startswith('#arch'):
                 self.arch = line.split('=')[1]
                 continue
+
             if line.startswith('['):
                 current_key = WineRegistryKey(key_def=line)
                 self.keys[current_key.name] = current_key
                 continue
+
             if current_key:
                 current_key.parse(line)
 
+    def render(self):
+        content = ""
+        content += "{}{}\n".format(self.version_header, self.version)
+        content += "{}{}\n\n".format(self.relative_to_header, self.relative_to)
+        content += "#arch={}\n".format(self.arch)
+        for key in self.keys:
+            content += "\n"
+            content += self.keys[key].render()
+        return content
+
     def query(self, keypath, value=None):
         key = self.keys.get(keypath)
         if key:
diff --git a/tests/test_registry.py b/tests/test_registry.py
index 6f72956d7..972198ef1 100644
--- a/tests/test_registry.py
+++ b/tests/test_registry.py
@@ -7,8 +7,8 @@ FIXTURES_PATH = os.path.join(os.path.dirname(__file__), 'fixtures')
 
 class TestWineRegistry(TestCase):
     def setUp(self):
-        registry_path = os.path.join(FIXTURES_PATH, 'user.reg')
-        self.registry = WineRegistry(registry_path)
+        self.registry_path = os.path.join(FIXTURES_PATH, 'user.reg')
+        self.registry = WineRegistry(self.registry_path)
 
     def test_can_load_registry(self):
         self.assertTrue(len(self.registry.keys) > 10)
@@ -48,3 +48,19 @@ class TestWineRegistry(TestCase):
         )
         key = self.registry.keys.get('Software/Wine/Fonts')
         self.assertEqual(key.render(), expected)
+
+    def test_render_user_reg(self):
+        content = self.registry.render()
+        with open(self.registry_path, 'r') as registry_file:
+            original_content = registry_file.read()
+        self.assertEqual(content, original_content)
+
+    def test_can_render_system_reg(self):
+        registry_path = os.path.join(FIXTURES_PATH, 'system.reg')
+        with open(registry_path, 'r') as registry_file:
+            original_content = registry_file.read()
+        system_reg = WineRegistry(registry_path)
+        content = system_reg.render()
+        with open(registry_path + '.new', 'w') as new_file:
+            new_file.write(content)
+        self.assertEqual(content, original_content)