Make gzip reproducible (#560)

Author: skearnes

ord_schema/message_helpers.py

@@ -14,6 +14,7 @@
 """Helper functions for constructing Protocol Buffer messages."""
 
 import enum
+import functools
 import gzip
 import os
 import re
@@ -770,21 +771,19 @@ def write_message(message: ord_schema.Message, filename: str):
         ValueError: if `filename` does not have the expected suffix.
     """
     if filename.endswith('.gz'):
-        this_open = gzip.open
+        # NOTE(kearnes): Set a constant mtime so that round-trips through gzip
+        # result in identical files.
+        this_open = functools.partial(gzip.GzipFile, mtime=1)
         _, extension = os.path.splitext('.'.join(filename.split('.')[:-1]))
     else:
         this_open = open
         _, extension = os.path.splitext(filename)
     output_format = MessageFormat(extension)
-    if output_format == MessageFormat.BINARY:
-        mode = 'wb'
-    else:
-        mode = 'wt'
-    with this_open(filename, mode) as f:
+    with this_open(filename, 'wb') as f:
         if output_format == MessageFormat.JSON:
-            f.write(json_format.MessageToJson(message))
+            f.write(json_format.MessageToJson(message).encode())
         elif output_format == MessageFormat.PBTXT:
-            f.write(text_format.MessageToString(message))
+            f.write(text_format.MessageToBytes(message))
         elif output_format == MessageFormat.BINARY:
             f.write(message.SerializeToString(deterministic=True))
 

ord_schema/message_helpers_test.py

@@ -15,6 +15,7 @@
 
 import os
 import tempfile
+import time
 
 from absl import flags
 from absl.testing import absltest
@@ -429,6 +430,7 @@ class LoadAndWriteMessageTest(parameterized.TestCase, absltest.TestCase):
 
     def setUp(self):
         super().setUp()
+        self.test_directory = self.create_tempdir()
         self.messages = [
             test_pb2.Scalar(int32_value=3, float_value=4.5),
             test_pb2.RepeatedScalar(values=[1.2, 3.4]),
@@ -448,6 +450,17 @@ def test_round_trip(self, suffix):
                     message,
                     message_helpers.load_message(f.name, type(message)))
 
+    def test_gzip_reproducibility(self):
+        filename = os.path.join(self.test_directory, 'test.pb.gz')
+        for message in self.messages:
+            message_helpers.write_message(message, filename)
+            with open(filename, 'rb') as f:
+                value = f.read()
+            time.sleep(1)
+            message_helpers.write_message(message, filename)
+            with open(filename, 'rb') as f:
+                self.assertEqual(f.read(), value)
+
     def test_bad_binary(self):
         with tempfile.NamedTemporaryFile(suffix='.pb') as f:
             message = test_pb2.RepeatedScalar(values=[1.2, 3.4])