Merge pull request #2018 from sydduckworth/block-checksum

Support not writing checksums

Author: sydduckworth

asdf/_asdf.py

@@ -768,7 +768,7 @@ def _tree_finalizer(tagged_tree):
             padding = util.calculate_padding(fd.tell(), pad_blocks, fd.block_size)
             fd.fast_forward(padding)
 
-    def _serial_write(self, fd, pad_blocks, include_block_index):
+    def _serial_write(self, fd, pad_blocks, include_block_index, write_checksums):
         with self._blocks.write_context(fd):
             # prep a tree for a writing
             tree = copy.copy(self._tree)
@@ -777,7 +777,7 @@ def _serial_write(self, fd, pad_blocks, include_block_index):
                 tree["history"] = copy.deepcopy(self._tree["history"])
 
             self._write_tree(tree, fd, pad_blocks)
-            self._blocks.write(pad_blocks, include_block_index)
+            self._blocks.write(pad_blocks, include_block_index, write_checksums)
 
     def update(
         self,
@@ -787,6 +787,7 @@ def update(
         pad_blocks=False,
         include_block_index=True,
         version=None,
+        write_checksums=True,
     ):
         """
         Update the file on disk in place.
@@ -839,6 +840,9 @@ def update(
         version : str, optional
             Update the ASDF core schemas version of this AsdfFile before
             writing.
+
+        write_checksums: bool, optional
+            Compute and write block checksums to the file.
         """
 
         with config_context() as config:
@@ -876,7 +880,12 @@ def update(
 
             def rewrite():
                 self._fd.seek(0)
-                self._serial_write(self._fd, pad_blocks, include_block_index)
+                self._serial_write(
+                    self._fd,
+                    pad_blocks,
+                    include_block_index,
+                    write_checksums,
+                )
                 self._fd.truncate()
                 if self._fd.can_memmap():
                     self._fd.close_memmap()
@@ -901,7 +910,7 @@ def rewrite():
                 new_tree_size = tree_fd.tell()
 
                 # update blocks
-                self._blocks.update(new_tree_size, pad_blocks, include_block_index)
+                self._blocks.update(new_tree_size, pad_blocks, include_block_index, write_checksums)
                 end_of_file = self._fd.tell()
 
             # now write the tree
@@ -925,6 +934,7 @@ def write_to(
         pad_blocks=False,
         include_block_index=True,
         version=None,
+        write_checksums=True,
     ):
         """
         Write the ASDF file to the given file-like object.
@@ -987,6 +997,9 @@ def write_to(
         version : str, optional
             Update the ASDF core schemas version of this AsdfFile before
             writing.
+
+        write_checksums: bool, optional
+            Compute and write block checksums to the file.
         """
         with config_context() as config:
             if all_array_storage is not NotSet:
@@ -1002,7 +1015,7 @@ def write_to(
 
             try:
                 with generic_io.get_file(fd, mode="w") as fd:
-                    self._serial_write(fd, pad_blocks, include_block_index)
+                    self._serial_write(fd, pad_blocks, include_block_index, write_checksums)
             finally:
                 if version is not None:
                     self.version = previous_version

asdf/_block/io.py

@@ -275,7 +275,9 @@ def callback():
     return offset, header, data_offset, data
 
 
-def generate_write_header(data, stream=False, compression_kwargs=None, padding=False, fs_block_size=1, **header_kwargs):
+def generate_write_header(
+    data, stream=False, compression_kwargs=None, padding=False, fs_block_size=1, write_checksum=True, **header_kwargs
+):
     """
     Generate a dict representation of a ASDF block header that can be
     used for writing a block.
@@ -307,6 +309,10 @@ def generate_write_header(data, stream=False, compression_kwargs=None, padding=F
         The filesystem block size. See the `asdf.util.calculate_padding`
         ``block_size`` argument for more details.
 
+    write_checksum: bool, optional
+        Compute and write the checksum of the block data.
+        If disabled then the checksum field is set to 0.
+
     **header_kwargs : dict, optional
         Block header settings that will be read, updated, and used
         to generate the binary block header representation by packing
@@ -353,7 +359,7 @@ def generate_write_header(data, stream=False, compression_kwargs=None, padding=F
         padding = util.calculate_padding(used_size, padding, fs_block_size)
         header_kwargs["allocated_size"] = header_kwargs.get("allocated_size", used_size + padding)
 
-    if stream:
+    if stream or not write_checksum:
         header_kwargs["checksum"] = b"\0" * 16
     elif buff is not None:
         header_kwargs["checksum"] = calculate_block_checksum(buff.getbuffer())
@@ -370,7 +376,9 @@ def generate_write_header(data, stream=False, compression_kwargs=None, padding=F
     return header_kwargs, buff, padding_bytes
 
 
-def write_block(fd, data, offset=None, stream=False, compression_kwargs=None, padding=False, **header_kwargs):
+def write_block(
+    fd, data, offset=None, stream=False, compression_kwargs=None, padding=False, write_checksum=True, **header_kwargs
+):
     """
     Write an ASDF block.
 
@@ -392,6 +400,10 @@ def write_block(fd, data, offset=None, stream=False, compression_kwargs=None, pa
     padding : bool, optional, default False
         Optionally pad the block data. See `generate_write_header`.
 
+    write_checksum: bool, optional
+        Compute and write the checksum of the block data.
+        If disabled then the checksum field is set to 0.
+
     **header_kwargs : dict
         Block header settings. See `generate_write_header`.
 
@@ -403,7 +415,7 @@ def write_block(fd, data, offset=None, stream=False, compression_kwargs=None, pa
         for writing.
     """
     header_dict, buff, padding_bytes = generate_write_header(
-        data, stream, compression_kwargs, padding, fd.block_size, **header_kwargs
+        data, stream, compression_kwargs, padding, fd.block_size, write_checksum, **header_kwargs
     )
     header_bytes = BLOCK_HEADER.pack(**header_dict)
 

asdf/_block/manager.py

@@ -364,7 +364,7 @@ def _clear_write(self):
         self._streamed_obj_keys = set()
         self._write_fd = None
 
-    def _write_external_blocks(self):
+    def _write_external_blocks(self, write_checksums):
         from asdf import AsdfFile
 
         if self._write_fd is None or self._write_fd.uri is None:
@@ -376,7 +376,7 @@ def _write_external_blocks(self):
             af = AsdfFile()
             with generic_io.get_file(uri, mode="w") as f:
                 af.write_to(f, include_block_index=False)
-                writer.write_blocks(f, [blk])
+                writer.write_blocks(f, [blk], write_checksums=write_checksums)
 
     def make_write_block(self, data, options, obj):
         """
@@ -541,7 +541,7 @@ def write_context(self, fd, copy_options=True):
             yield
         self._clear_write()
 
-    def write(self, pad_blocks, include_block_index):
+    def write(self, pad_blocks, include_block_index, write_checksums):
         """
         Write blocks that were set up during the current
         `write_context`.
@@ -559,6 +559,9 @@ def write(self, pad_blocks, include_block_index):
             If a streamed_block is provided (or the file is not
             seekable) no block index will be written.
 
+        write_checksums : bool
+            Compute and write checksums for each block.
+
         Raises
         ------
         OSError
@@ -574,11 +577,12 @@ def write(self, pad_blocks, include_block_index):
                 pad_blocks,
                 streamed_block=self._streamed_write_block,
                 write_index=include_block_index,
+                write_checksums=write_checksums,
             )
         if len(self._external_write_blocks):
-            self._write_external_blocks()
+            self._write_external_blocks(write_checksums=write_checksums)
 
-    def update(self, new_tree_size, pad_blocks, include_block_index):
+    def update(self, new_tree_size, pad_blocks, include_block_index, write_checksums):
         """
         Perform an update-in-place of ASDF blocks set up during
         a `write_context`.
@@ -596,11 +600,13 @@ def update(self, new_tree_size, pad_blocks, include_block_index):
             a number of padding bytes based off a ratio of the data
             size.
 
-        include_block_index : bool
+        include_block_index : bool or None
             If True, include a block index at the end of the file.
             If a streamed_block is provided (or the file is not
             seekable) no block index will be written.
 
+        write_checksums: bool, optional
+            Compute and write block checksums to the file.
 
         Raises
         ------
@@ -626,7 +632,7 @@ def update(self, new_tree_size, pad_blocks, include_block_index):
             )
 
         if len(self._external_write_blocks):
-            self._write_external_blocks()
+            self._write_external_blocks(write_checksums=write_checksums)
 
         # do we have any blocks to write?
         if len(self._write_blocks) or self._streamed_write_block:
@@ -637,6 +643,7 @@ def update(self, new_tree_size, pad_blocks, include_block_index):
                 pad_blocks,
                 streamed_block=self._streamed_write_block,
                 write_index=False,  # don't write an index as we will modify the offsets
+                write_checksums=write_checksums,
             )
             new_block_end = self._write_fd.tell()
 

asdf/_block/writer.py

@@ -29,7 +29,7 @@ def data_bytes(self):
         return np.ndarray(0, np.uint8)
 
 
-def write_blocks(fd, blocks, padding=False, streamed_block=None, write_index=True):
+def write_blocks(fd, blocks, padding=False, streamed_block=None, write_index=True, write_checksums=True):
     """
     Write a list of WriteBlocks to a file
 
@@ -60,6 +60,9 @@ def write_blocks(fd, blocks, padding=False, streamed_block=None, write_index=Tru
         If a streamed_block is provided (or the file is not
         seekable) no block index will be written.
 
+    write_checksums: bool, optional
+        Compute and write block checksums to the file.
+
     Returns
     -------
     offsets : list of int
@@ -101,12 +104,13 @@ def tell():
                 compression_kwargs=blk.compression_kwargs,
                 padding=padding,
                 compression=blk.compression,
+                write_checksum=write_checksums,
             )
         )
     if streamed_block is not None:
         offsets.append(tell())
         fd.write(constants.BLOCK_MAGIC)
-        headers.append(bio.write_block(fd, streamed_block.data_bytes, stream=True))
+        headers.append(bio.write_block(fd, streamed_block.data_bytes, stream=True, write_checksum=write_checksums))
 
     # os.pipe on windows returns a file-like object
     # that reports as seekable but tell always returns 0

asdf/_dump.py

@@ -17,6 +17,7 @@ def dump(
     compression_kwargs=NotSet,
     pad_blocks=False,
     custom_schema=None,
+    write_checksums=True,
 ):
     """
     Write a tree to an ASDF file.
@@ -54,6 +55,9 @@ def dump(
         validation pass. This can be used to ensure that particular ASDF
         files follow custom conventions beyond those enforced by the
         specification.
+
+    write_checksums: bool, optional
+        Compute and write block checksums to the file.
     """
     AsdfFile(tree, custom_schema=custom_schema, extensions=extensions).write_to(
         fp,
@@ -62,6 +66,7 @@ def dump(
         all_array_compression=all_array_compression,
         compression_kwargs=compression_kwargs,
         pad_blocks=pad_blocks,
+        write_checksums=write_checksums,
     )
 
 
@@ -75,6 +80,7 @@ def dumps(
     compression_kwargs=NotSet,
     pad_blocks=False,
     custom_schema=None,
+    write_checksums=True,
 ):
     """
     Write tree to a string.
@@ -110,6 +116,9 @@ def dumps(
         files follow custom conventions beyond those enforced by the
         specification.
 
+    write_checksums: bool, optional
+        Compute and write block checksums to the file.
+
     Returns
     -------
     str
@@ -126,6 +135,7 @@ def dumps(
         compression_kwargs=compression_kwargs,
         pad_blocks=pad_blocks,
         custom_schema=custom_schema,
+        write_checksums=write_checksums,
     )
     return buff.getvalue()
 

asdf/_tests/_block/test_io.py

@@ -18,7 +18,8 @@ def test_checksum(tmp_path):
     # check that when written, a block generates the correct checksum
     path = tmp_path / "test"
     with generic_io.get_file(path, mode="w") as fd:
-        bio.write_block(fd, my_array.view(dtype="uint8"))
+        # check that no warnings are raised when writing checksums without compression
+        bio.write_block(fd, my_array.view(dtype="uint8"), write_checksum=True)
     with generic_io.get_file(path, mode="r") as fd:
         _, header, _, _ = bio.read_block(fd, True)
     assert header["checksum"] == target_checksum

asdf/_tests/_block/test_manager.py

@@ -58,7 +58,7 @@ def test_write_no_uri(tmp_path):
     with asdf.open(fn) as af:
         m = af._blocks
         with pytest.raises(ValueError, match=r"Can't write external blocks.*"):
-            m._write_external_blocks()
+            m._write_external_blocks(write_checksums=False)
 
 
 def test_write_outside_context(tmp_path):
@@ -67,7 +67,7 @@ def test_write_outside_context(tmp_path):
     with asdf.open(fn) as af:
         m = af._blocks
         with pytest.raises(OSError, match=r"write called outside of valid write_context"):
-            m.write(False, False)
+            m.write(pad_blocks=False, include_block_index=False, write_checksums=False)
 
 
 def test_update_outside_context(tmp_path):
@@ -76,7 +76,7 @@ def test_update_outside_context(tmp_path):
     with asdf.open(fn) as af:
         m = af._blocks
         with pytest.raises(OSError, match=r"update called outside of valid write_context"):
-            m.update(0, False, False)
+            m.update(0, pad_blocks=False, include_block_index=False, write_checksums=False)
 
 
 def test_input_compression(tmp_path):