Support for adding custom blocks on write

mikeio/dataset/_dataset.py

@@ -1834,7 +1834,8 @@ def to_dfs(self, filename: str | Path, **kwargs: Any) -> None:
                 raise ValueError("Cannot write Dataset with no geometry to file!")
         elif isinstance(self.geometry, Grid2D):
             self._validate_extension(filename, ".dfs2")
-            self._to_dfs2(filename)
+            custom_blocks = kwargs.get("custom_blocks", None)
+            self._to_dfs2(filename, custom_blocks=custom_blocks)
         elif isinstance(self.geometry, Grid3D):
             self._validate_extension(filename, ".dfs3")
             self._to_dfs3(filename)
@@ -1864,11 +1865,15 @@ def _to_dfs0(self, filename: str | Path, **kwargs: Any) -> None:
 
         _write_dfs0(filename, self, dtype=dtype)
 
-    def _to_dfs2(self, filename: str | Path) -> None:
+    def _to_dfs2(
+        self,
+        filename: str | Path,
+        custom_blocks: Mapping[str, np.ndarray] | None = None,
+    ) -> None:
         # assumes Grid2D geometry
         from ..dfs._dfs2 import write_dfs2
 
-        write_dfs2(filename, self)
+        write_dfs2(filename, self, custom_blocks=custom_blocks)
 
     def _to_dfs3(self, filename: str | Path) -> None:
         # assumes Grid3D geometry

mikeio/dfs/_dfs2.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, List, Literal, Tuple
+from typing import Any, List, Literal, Mapping, Tuple
 from collections.abc import Sequence
 
 import numpy as np
@@ -27,14 +27,42 @@
 from ..spatial import Grid2D
 
 
-def write_dfs2(filename: str | Path, ds: Dataset, title: str = "") -> None:
-    dfs = _write_dfs2_header(filename, ds, title)
+def write_dfs2(
+    filename: str | Path,
+    ds: Dataset,
+    title: str = "",
+    custom_blocks: Mapping[str, np.ndarray] | None = None,
+) -> None:
+    """Write a dfs2 file
+
+    Parameters
+    ----------
+    filename: str or Path
+        File name to write
+    ds: Dataset
+        The dataset to write
+    title: str, optional
+        Title of the file
+    custom_blocks: dict, optional
+        Custom blocks to add to the file
+    """
+
+    dfs = _write_dfs2_header(filename, ds, title, custom_blocks=custom_blocks)
     write_dfs_data(dfs=dfs, ds=ds, n_spatial_dims=2)
 
 
-def _write_dfs2_header(filename: str | Path, ds: Dataset, title: str = "") -> DfsFile:
+def _write_dfs2_header(
+    filename: str | Path,
+    ds: Dataset,
+    title: str = "",
+    custom_blocks: Mapping[str, np.ndarray] | None = None,
+) -> DfsFile:
     builder = DfsBuilder.Create(title, "mikeio", __dfs_version__)
     builder.SetDataType(0)
+    if custom_blocks is not None:
+        factory = DfsFactory()
+        for key, value in custom_blocks.items():
+            builder.AddCustomBlock(factory.CreateCustomBlock(key, value))
 
     geometry: Grid2D = ds.geometry
 

tests/test_dfs2.py

@@ -825,6 +825,18 @@ def test_read_dfs2_static_dt_zero():
     assert "time" not in ds2.dims
 
 
+
+def test_add_custom_block_to_dfs2(tmp_path):
+    ds = mikeio.read("tests/testdata/gebco_sound.dfs2")
+    fp = tmp_path / "gebco_sound_with_custom_block.dfs2"
+    ds.to_dfs(
+        fp,
+        custom_blocks={"M21_Misc": np.array([327, 0.2, -900, 10, 0, 0, 0], np.float32)},
+    )
+
+    ds2 = mikeio.open(fp)
+    assert ds2._dfs.FileInfo.CustomBlocks[0].Name == "M21_Misc"
+
 def test_write_read_local_coordinates(tmp_path):
     da = mikeio.DataArray(
         np.array([[1, 2, 3], [4, 5, 6]]),
@@ -864,3 +876,4 @@ def test_to_xarray():
     xr_da = da.to_xarray()
     assert xr_da.x[0] == pytest.approx(0.25)
     assert xr_da.y[0] == pytest.approx(0.25)
+