Return xr_dataset from create_empty

Author: dmitriyrepin

src/mdio/api/create.py

@@ -30,10 +30,10 @@
 def create_empty(  # noqa PLR0913
     mdio_template: AbstractDatasetTemplate | str,
     dimensions: list[Dimension],
-    output_path: UPath | Path | str,
+    output_path: UPath | Path | str | None,
     headers: HeaderSpec | None = None,
     overwrite: bool = False,
-) -> None:
+) -> xr_Dataset:
     """A function that creates an empty MDIO v1 file with known dimensions.
 
     Args:
@@ -73,19 +73,24 @@ def create_empty(  # noqa PLR0913
     # Populate coordinates using the grid
     # For empty datasets, we only populate dimension coordinates
     drop_vars_delayed = []
-    dataset, drop_vars_delayed = populate_dim_coordinates(xr_dataset, grid, drop_vars_delayed=drop_vars_delayed)
+    xr_dataset, drop_vars_delayed = populate_dim_coordinates(xr_dataset, grid, drop_vars_delayed=drop_vars_delayed)
 
     if headers:
         # Since the headers were provided, the user wants to export to SEG-Y
         # Add a dummy segy_file_header variable used to export to SEG-Y
-        dataset["segy_file_header"] = ((), "")
+        xr_dataset["segy_file_header"] = ((), "")
 
     # Create the Zarr store with the correct structure but with empty arrays
-    to_mdio(dataset, output_path=output_path, mode="w", compute=False)
+    if output_path is not None:
+        to_mdio(xr_dataset, output_path=output_path, mode="w", compute=False)
 
     # Write the dimension coordinates and trace mask
-    meta_ds = dataset[drop_vars_delayed + ["trace_mask"]]
-    to_mdio(meta_ds, output_path=output_path, mode="r+", compute=True)
+    xr_dataset = xr_dataset[drop_vars_delayed + ["trace_mask"]]
+    
+    if output_path is not None:
+        to_mdio(xr_dataset, output_path=output_path, mode="r+", compute=True)
+
+    return xr_dataset
 
 
 def create_empty_like(  # noqa PLR0913

tests/integration/test_z_create_empty.py

@@ -104,13 +104,14 @@ def _create_empty_mdio(cls, create_headers: bool, output_path: Path, overwrite:
         # The HeaderSpec can be either standard or customized.
         headers = get_teapot_segy_spec().trace.header if create_headers else None
         # Create an empty MDIO v1 metric post-stack 3D time velocity dataset
-        create_empty(
+        xr_dataset = create_empty(
             mdio_template=PostStack3DVelocityTemplate(data_domain="time", is_metric=True),
             dimensions=dims,
             output_path=output_path,
             headers=headers,
             overwrite=overwrite,
         )
+        return xr_dataset
 
     @classmethod
     def validate_teapod_dataset_metadata(cls, ds: xr_Dataset, is_velocity: bool) -> None:
@@ -203,7 +204,8 @@ def mdio_with_headers(self, empty_mdio_dir: Path) -> Path:
         and shared across all test methods in the class.
         """
         empty_mdio: Path = empty_mdio_dir / "mdio_with_headers.mdio"
-        self._create_empty_mdio(create_headers=True, output_path=empty_mdio)
+        xr_dataset = self._create_empty_mdio(create_headers=True, output_path=empty_mdio)
+        assert xr_dataset is not None
         return empty_mdio
 
     @pytest.fixture(scope="class")
@@ -214,7 +216,8 @@ def mdio_no_headers(self, empty_mdio_dir: Path) -> Path:
         and shared across all test methods in the class.
         """
         empty_mdio: Path = empty_mdio_dir / "mdio_no_headers.mdio"
-        self._create_empty_mdio(create_headers=False, output_path=empty_mdio)
+        xr_dataset = self._create_empty_mdio(create_headers=False, output_path=empty_mdio)
+        assert xr_dataset is not None
         return empty_mdio
 
     def test_dataset_metadata(self, mdio_with_headers: Path) -> None:
@@ -251,7 +254,8 @@ def test_overwrite_behavior(self, empty_mdio_dir: Path) -> None:
             self._create_empty_mdio(create_headers=True, output_path=empty_mdio, overwrite=False)
 
         # Third call: Create MDIO with overwrite=True - should succeed and overwrite garbage
-        self._create_empty_mdio(create_headers=True, output_path=empty_mdio, overwrite=True)
+        xr_dataset = self._create_empty_mdio(create_headers=True, output_path=empty_mdio, overwrite=True)
+        assert xr_dataset is not None
 
         # Validate that the MDIO file can be loaded correctly using the helper function
         ds = open_mdio(empty_mdio)
@@ -398,6 +402,8 @@ def test_create_empty_like(self, teapot_mdio_tmp: Path, empty_mdio_dir: Path) ->
             keep_coordinates=True,
             overwrite=True,
         )
+        assert ds is not None
+
         self.validate_teapod_dataset_metadata(ds, is_velocity=False)
         header_dtype = get_teapot_segy_spec().trace.header.dtype
         self.validate_teapod_dataset_variables(ds, header_dtype=header_dtype, is_velocity=False)