locale_database: Use context manager interface to update source files

Use context manager interface (with statement) to atomically update source
files. This ensures that all files are properly closed and the temporary
file is removed even in case of errors.

Task-number: QTBUG-83488
Pick-to: 6.2
Change-Id: I18cd96f1d03e467dea6212f6576a41e65f414ce1
Reviewed-by: Edward Welbourne <edward.welbourne@qt.io>
Reviewed-by: Cristian Maureira-Fredes <cristian.maureira-fredes@qt.io>

3 files changed, 134 insertions, 129 deletions

diff --git a/util/locale_database/cldr2qtimezone.py b/util/locale_database/cldr2qtimezone.py
index bd9b050e1c..dec350fa8e 100755
--- a/util/locale_database/cldr2qtimezone.py
+++ b/util/locale_database/cldr2qtimezone.py
@@ -369,20 +369,14 @@ def main(out, err):
         return 1
 
     out.write('Input file parsed, now writing data\n')
-    try:
-        writer = ZoneIdWriter(dataFilePath, qtPath)
-    except IOError as e:
-        err.write(f'Failed to open files to transcribe: {e}')
-        return 1
 
     try:
-        writer.write(version, defaults, winIds)
-    except Error as e:
-        writer.cleanup()
-        err.write(f'\nError in Windows ID data: {e}\n')
+        with ZoneIdWriter(dataFilePath, qtPath) as writer:
+            writer.write(version, defaults, winIds)
+    except Exception as e:
+        err.write(f'\nError while updating timezone data: {e}\n')
         return 1
 
-    writer.close()
     out.write(f'Data generation completed, please check the new file at {dataFilePath}\n')
     return 0
 
diff --git a/util/locale_database/localetools.py b/util/locale_database/localetools.py
index 37b39a07e0..ee6abd5593 100644
--- a/util/locale_database/localetools.py
+++ b/util/locale_database/localetools.py
@@ -37,6 +37,7 @@ Classes:
   SourceFileEditor -- adds standard prelude and tail handling to Transcriber.
 """
 
+from contextlib import ExitStack, contextmanager
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 
@@ -69,45 +70,98 @@ def wrap_list(lst):
             yield head
     return ",\n".join(", ".join(x) for x in split(lst, 20))
 
-class Transcriber (object):
-    """Helper class to facilitate rewriting source files.
 
-    This class takes care of the temporary file manipulation. Derived
-    classes need to implement transcribing of the content, with
+@contextmanager
+def AtomicRenameTemporaryFile(originalLocation: Path, *, prefix: str, dir: Path):
+    """Context manager for safe file update via a temporary file.
+
+    Accepts path to the file to be updated. Yields a temporary file to the user
+    code, open for writing.
+
+    On success closes the temporary file and moves its content to the original
+    location. On error, removes temporary file, without disturbing the original.
+    """
+    tempFile = NamedTemporaryFile('w', prefix=prefix, dir=dir, delete=False)
+    try:
+        yield tempFile
+        tempFile.close()
+        # Move the modified file to the original location
+        Path(tempFile.name).rename(originalLocation)
+    except Exception:
+        # delete the temporary file in case of error
+        tempFile.close()
+        Path(tempFile.name).unlink()
+        raise
+
+
+class Transcriber:
+    """Context manager base-class to manage source file rewrites.
+
+    Derived classes need to implement transcribing of the content, with
     whatever modifications they may want.  Members reader and writer
     are exposed; use writer.write() to output to the new file; use
     reader.readline() or iterate reader to read the original.
 
-    Callers should call close() on success or cleanup() on failure (to
-    clear away the temporary file).
+    This class is intended to be used as context manager only (inside a
+    `with` statement).
+
+    Reimplement onEnter() to write any preamble the file may have,
+    onExit() to write any tail. The body of the with statement takes
+    care of anything in between, using methods provided by derived classes.
+
+    The data is written to a temporary file first. The temporary file data
+    is then moved to the original location if there were no errors. Otherwise
+    the temporary file is removed and the original is left unchanged.
     """
     def __init__(self, path: Path, temp_dir: Path):
-        # Open the old file
-        self.reader = open(path)
         self.path = path
-        # Create a temp file to write the new data into
-        self.writer = NamedTemporaryFile('w', prefix=path.name, dir=temp_dir, delete=False)
+        self.tempDir = temp_dir
 
-    def close(self) -> None:
-        self.reader.close()
-        self.writer.close()
+    def onEnter(self) -> None:
+        """
+        Called before transferring control to user code.
 
-        # Move the modified file to the original location
-        self.path.unlink()
-        Path(self.writer.name).rename(self.path)
+        This function can be overridden in derived classes to perform actions
+        before transferring control to the user code.
+
+        The default implementation does nothing.
+        """
+        pass
 
-        self.reader = self.writer = None
+    def onExit(self) -> None:
+        """
+        Called after return from user code.
 
-    def cleanup(self) -> None:
-        if self.reader:
-            self.reader.close()
-            self.reader = None
+        This function can be overridden in derived classes to perform actions
+        after successful return from user code.
 
-        if self.writer:
-            self.writer.close()
-            # Remove temp-file:
-            Path(self.writer.name).unlink(missing_ok=True)
-            self.writer = None
+        The default implementation does nothing.
+        """
+        pass
+
+    def __enter__(self):
+        with ExitStack() as resources:
+            # Create a temp file to write the new data into
+            self.writer = resources.enter_context(
+                AtomicRenameTemporaryFile(self.path, prefix=self.path.name, dir=self.tempDir))
+            # Open the old file
+            self.reader = resources.enter_context(open(self.path))
+
+            self.onEnter()
+
+            # Prevent resources from being closed on normal return from this
+            # method and make them available inside __exit__():
+            self.__resources = resources.pop_all()
+            return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            with self.__resources:
+               self.onExit()
+        else:
+            self.__resources.__exit__(exc_type, exc_value, traceback)
+
+        return False
 
 
 class SourceFileEditor (Transcriber):
@@ -121,41 +175,27 @@ class SourceFileEditor (Transcriber):
     the new version to replace it.
 
     This class takes care of transcribing the parts before and after
-    the generated content; on creation, an instance will copy the
-    preamble up to the start marker; its close() will skip over the
-    original's generated content and resume transcribing with the end
-    marker. Derived classes need only implement the generation of the
-    content in between.
-
-    Callers should call close() on success or cleanup() on failure (to
-    clear away the temporary file); see Transcriber.
+    the generated content; on entering the context, an instance will
+    copy the preamble up to the start marker; on exit from the context
+    it will skip over the original's generated content and resume
+    transcribing with the end marker.
+
+    This class is only intended to be used as a context manager:
+    see Transcriber. Derived classes implement suitable methods for use in
+    the body of the with statement, using self.writer to rewrite the part
+    of the file between the start and end markers.
     """
-    def __init__(self, path: Path, temp_dir: Path):
-        """Set up the source file editor.
-
-        Requires two arguments: the path to the source file to be read
-        and, on success, replaced with a new version; and the
-        directory in which to store the temporary file during the
-        rewrite."""
-        super().__init__(path, temp_dir)
-        self.__copyPrelude()
-
-    def close(self):
-        self.__copyTail()
-        super().close()
-
-    # Implementation details:
     GENERATED_BLOCK_START = '// GENERATED PART STARTS HERE'
     GENERATED_BLOCK_END = '// GENERATED PART ENDS HERE'
 
-    def __copyPrelude(self):
+    def onEnter(self) -> None:
         # Copy over the first non-generated section to the new file
         for line in self.reader:
             self.writer.write(line)
             if line.strip() == self.GENERATED_BLOCK_START:
                 break
 
-    def __copyTail(self):
+    def onExit(self) -> None:
         # Skip through the old generated data in the old file
         for line in self.reader:
             if line.strip() == self.GENERATED_BLOCK_END:
diff --git a/util/locale_database/qlocalexml2cpp.py b/util/locale_database/qlocalexml2cpp.py
index 9866bd6ea7..7ac7945cf8 100755
--- a/util/locale_database/qlocalexml2cpp.py
+++ b/util/locale_database/qlocalexml2cpp.py
@@ -135,10 +135,14 @@ def currencyIsoCodeData(s):
 class LocaleSourceEditor (SourceFileEditor):
     def __init__(self, path: Path, temp: Path, version: str):
         super().__init__(path, temp)
+        self.version = version
+
+    def onEnter(self) -> None:
+        super().onEnter()
         self.writer.write(f"""
 /*
     This part of the file was generated on {datetime.date.today()} from the
-    Common Locale Data Repository v{version}
+    Common Locale Data Repository v{self.version}
 
     http://www.unicode.org/cldr/
 
@@ -535,89 +539,56 @@ def main(out, err):
     locale_keys = sorted(locale_map.keys(), key=LocaleKeySorter(reader.defaultMap()))
 
     try:
-        writer = LocaleDataWriter(qtsrcdir.joinpath('src/corelib/text/qlocale_data_p.h'),
-                                  qtsrcdir, reader.cldrVersion)
-    except IOError as e:
-        err.write(f'Failed to open files to transcribe locale data: {e}')
-        return 1
-
-    try:
-        writer.likelySubtags(reader.likelyMap())
-        writer.localeIndex(reader.languageIndices(tuple(k[0] for k in locale_map)))
-        writer.localeData(locale_map, locale_keys)
-        writer.writer.write('\n')
-        writer.languageNames(reader.languages)
-        writer.scriptNames(reader.scripts)
-        writer.territoryNames(reader.territories)
-        # TODO: merge the next three into the previous three
-        writer.languageCodes(reader.languages)
-        writer.scriptCodes(reader.scripts)
-        writer.territoryCodes(reader.territories)
-    except Error as e:
-        writer.cleanup()
+        with LocaleDataWriter(qtsrcdir.joinpath('src/corelib/text/qlocale_data_p.h'),
+                              qtsrcdir, reader.cldrVersion) as writer:
+            writer.likelySubtags(reader.likelyMap())
+            writer.localeIndex(reader.languageIndices(tuple(k[0] for k in locale_map)))
+            writer.localeData(locale_map, locale_keys)
+            writer.writer.write('\n')
+            writer.languageNames(reader.languages)
+            writer.scriptNames(reader.scripts)
+            writer.territoryNames(reader.territories)
+            # TODO: merge the next three into the previous three
+            writer.languageCodes(reader.languages)
+            writer.scriptCodes(reader.scripts)
+            writer.territoryCodes(reader.territories)
+    except Exception as e:
         err.write(f'\nError updating locale data: {e}\n')
         return 1
 
-    writer.close()
-
     # Generate calendar data
     for calendar, stem in calendars.items():
         try:
-            writer = CalendarDataWriter(
-                qtsrcdir.joinpath(f'src/corelib/time/q{stem}calendar_data_p.h'),
-                qtsrcdir, reader.cldrVersion)
-        except IOError as e:
-            err.write(f'Failed to open files to transcribe {calendar} data {e}')
-            return 1
-
-        try:
-            writer.write(calendar, locale_map, locale_keys)
-        except Error as e:
-            writer.cleanup()
+            with CalendarDataWriter(
+                    qtsrcdir.joinpath(f'src/corelib/time/q{stem}calendar_data_p.h'),
+                    qtsrcdir, reader.cldrVersion) as writer:
+                writer.write(calendar, locale_map, locale_keys)
+        except Exception as e:
             err.write(f'\nError updating {calendar} locale data: {e}\n')
-            return 1
-
-        writer.close()
 
     # qlocale.h
     try:
-        writer = LocaleHeaderWriter(qtsrcdir.joinpath('src/corelib/text/qlocale.h'),
-                                    qtsrcdir, reader.dupes)
-    except IOError as e:
-        err.write(f'Failed to open files to transcribe qlocale.h: {e}')
-        return 1
-
-    try:
-        writer.languages(reader.languages)
-        writer.scripts(reader.scripts)
-        writer.territories(reader.territories)
-    except Error as e:
-        writer.cleanup()
+        with LocaleHeaderWriter(qtsrcdir.joinpath('src/corelib/text/qlocale.h'),
+                                qtsrcdir, reader.dupes) as writer:
+            writer.languages(reader.languages)
+            writer.scripts(reader.scripts)
+            writer.territories(reader.territories)
+    except Exception as e:
         err.write(f'\nError updating qlocale.h: {e}\n')
-        return 1
-
-    writer.close()
 
     # qlocale.qdoc
     try:
-        writer = Transcriber(qtsrcdir.joinpath('src/corelib/text/qlocale.qdoc'), qtsrcdir)
-    except IOError as e:
-        err.write(f'Failed to open files to transcribe qlocale.qdoc: {e}')
-        return 1
-
-    DOCSTRING = "    QLocale's data is based on Common Locale Data Repository "
-    try:
-        for line in writer.reader:
-            if DOCSTRING in line:
-                writer.writer.write(f'{DOCSTRING}v{reader.cldrVersion}.\n')
-            else:
-                writer.writer.write(line)
-    except Error as e:
-        writer.cleanup()
-        err.write(f'\nError updating qlocale.qdoc: {e}\n')
+        with Transcriber(qtsrcdir.joinpath('src/corelib/text/qlocale.qdoc'), qtsrcdir) as qdoc:
+            DOCSTRING = "    QLocale's data is based on Common Locale Data Repository "
+            for line in qdoc.reader:
+                if DOCSTRING in line:
+                    qdoc.writer.write(f'{DOCSTRING}v{reader.cldrVersion}.\n')
+                else:
+                    qdoc.writer.write(line)
+    except Exception as e:
+        err.write(f'\nError updating qlocale.h: {e}\n')
         return 1
 
-    writer.close()
     return 0
 
 if __name__ == "__main__":