Merge pull request #2198 from swahtz/fix/nanovdb_python_write_grids_drop

nanovdb/python: fix writeGrids dropping all but the last grid

Author: swahtz

nanovdb/nanovdb/python/PyIO.cc

@@ -69,8 +69,15 @@ template<typename BufferT> nb::list readGrids(const std::string& fileName, int v
 
 template<typename BufferT> void writeGrids(const std::string& fileName, nb::list handles, io::Codec codec, int verbose)
 {
+    std::ofstream os(fileName, std::ios::out | std::ios::binary | std::ios::trunc);
+    if (!os.is_open()) {
+        throw std::ios_base::failure("Unable to open file named \"" + fileName + "\" for output");
+    }
     for (size_t i = 0; i < handles.size(); ++i) {
-        nanovdb::io::writeGrid(fileName, nb::cast<const GridHandle<BufferT>&>(handles[i]), codec, verbose);
+        nanovdb::io::writeGrid(os, nb::cast<const GridHandle<BufferT>&>(handles[i]), codec);
+    }
+    if (verbose) {
+        std::cout << "Wrote " << handles.size() << " nanovdb::Grid(s) to file named \"" << fileName << "\"" << std::endl;
     }
 }
 

nanovdb/nanovdb/python/test/TestNanoVDB.py

@@ -343,6 +343,17 @@ def test_list_to_vector(self):
         dstFile.close()
         try:
             nanovdb.io.writeGrids(dstFile.name, handles)
+            # Verify that both grids actually made it into the file. Previously
+            # writeGrids re-opened the file per handle with std::ios::trunc, so
+            # only the final grid was retained on disk.
+            metadata = nanovdb.io.readGridMetaData(dstFile.name)
+            self.assertEqual(len(metadata), len(handles))
+            readHandles = nanovdb.io.readGrids(dstFile.name)
+            self.assertEqual(len(readHandles), len(handles))
+            for readHandle in readHandles:
+                self.assertEqual(readHandle.gridCount(), 1)
+                self.assertEqual(readHandle.gridType(0), nanovdb.GridType.Double)
+                self.assertIsNotNone(readHandle.doubleGrid())
         finally:
             os.unlink(dstFile.name)
 
@@ -510,6 +521,22 @@ def test_read_write_grids(self):
         except RuntimeError:
             print("ZIP compression codec not supported. Skipping...")
 
+    def test_device_write_grids_multi(self):
+        # Regression test for deviceWriteGrids: all grids in the list must end
+        # up in the output file, not just the last one.
+        handle = nanovdb.tools.cuda.createLevelSetSphere(
+            nanovdb.GridType.Float, name=self.gridName
+        )
+        handles = [handle, handle]
+        nanovdb.io.deviceWriteGrids(self.dstFile.name, handles)
+        metadata = nanovdb.io.readGridMetaData(self.dstFile.name)
+        self.assertEqual(len(metadata), len(handles))
+        readHandles = nanovdb.io.deviceReadGrids(self.dstFile.name)
+        self.assertEqual(len(readHandles), len(handles))
+        for readHandle in readHandles:
+            self.assertEqual(readHandle.gridCount(), 1)
+            self.assertEqual(readHandle.gridType(0), nanovdb.GridType.Float)
+
 
 @unittest.skipIf(
     not nanovdb.isCudaAvailable(), "nanovdb module was compiled without CUDA support"