BUG: Force invocation of closeZip from our code

dzenanz · Sep 1, 2023 · d3540bb · d3540bb
1 parent d92832b
commit d3540bb
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 7 deletions.
diff --git a/include/itkOMEZarrNGFFImageIO.h b/include/itkOMEZarrNGFFImageIO.h
@@ -100,6 +100,14 @@ class IOOMEZarrNGFF_EXPORT OMEZarrNGFFImageIO : public ImageIOBase
     size_t size;
   };
 
+  /** Construct a "magic" file name from the provided bufferInfo. */
+  static std::string
+  MakeMemoryFileName(const itk::OMEZarrNGFFImageIO::BufferInfo & bufferInfo)
+  {
+    size_t bufferInfoAddress = reinterpret_cast<size_t>(&bufferInfo);
+    return std::to_string(bufferInfoAddress) + ".memory";
+  }
+
   /*-------- This part of the interfaces deals with reading data. ----- */
 
   /** Determine the file type. Returns true if this ImageIO can read the
@@ -231,6 +239,11 @@ class IOOMEZarrNGFF_EXPORT OMEZarrNGFFImageIO : public ImageIOBase
   int                m_TimeIndex = INVALID_INDEX;
   int                m_ChannelIndex = INVALID_INDEX;
   AxesCollectionType m_StoreAxes;
+
+  const static unsigned m_EmptyZipSize = 22;
+  char                  m_EmptyZip[m_EmptyZipSize] = "PK\x05\x06"; // the rest is filled with zeroes
+  const BufferInfo      m_EmptyZipBufferInfo{ m_EmptyZip, m_EmptyZipSize };
+  const std::string     m_EmptyZipFileName = MakeMemoryFileName(m_EmptyZipBufferInfo);
 };
 } // end namespace itk
 

diff --git a/src/itkOMEZarrNGFFImageIO.cxx b/src/itkOMEZarrNGFFImageIO.cxx
@@ -635,7 +635,8 @@ OMEZarrNGFFImageIO::Read(void * buffer)
   }
 
   if (false)
-  {}
+  {
+  }
   READ_ELEMENT_IF(float)
   READ_ELEMENT_IF(int8_t)
   READ_ELEMENT_IF(uint8_t)
@@ -757,7 +758,10 @@ OMEZarrNGFFImageIO::WriteImageInformation()
 void
 OMEZarrNGFFImageIO::Write(const void * buffer)
 {
-  tsContext = tensorstore::Context::Default(); // start with clean zip handles
+  if (m_FileName.substr(m_FileName.size() - 4) == ".zip" || m_FileName.substr(m_FileName.size() - 7) == ".memory")
+  {
+    tsContext = tensorstore::Context::Default(); // start with clean zip handles
+  }
   this->WriteImageInformation();
 
   if (itkToTensorstoreComponentType(this->GetComponentType()) == tensorstore::dtype_v<void>)
@@ -791,7 +795,8 @@ OMEZarrNGFFImageIO::Write(const void * buffer)
   std::string driver = getKVstoreDriver(this->GetFileName());
 
   if (false) // start with a plain "if"
-  {}         // so element statements can all be "else if"
+  {
+  } // so element statements can all be "else if"
   ELEMENT_WRITE(int8_t)
   ELEMENT_WRITE(uint8_t)
   ELEMENT_WRITE(int16_t)
@@ -807,8 +812,13 @@ OMEZarrNGFFImageIO::Write(const void * buffer)
     itkExceptionMacro("Unsupported component type: " << GetComponentTypeAsString(this->GetComponentType()));
   }
 
-  // Create a new context to close the open zip handles
-  tsContext = tensorstore::Context::Default();
+  if (m_FileName.substr(m_FileName.size() - 4) == ".zip" || m_FileName.substr(m_FileName.size() - 7) == ".memory")
+  {
+    // Attempt to read a non-existent file from the in-memory zip to close the current one
+    nlohmann::json temp;
+    bool           wasRead = jsonRead(m_EmptyZipFileName + "/non-existent.json", temp, "zip");
+    assert(wasRead == false);
+  }
 }
 
 

diff --git a/test/itkOMEZarrNGFFInMemoryTest.cxx b/test/itkOMEZarrNGFFInMemoryTest.cxx
@@ -57,8 +57,7 @@ doTest(const char * inputFileName, const char * outputFileName)
 
   itk::OMEZarrNGFFImageIO::BufferInfo bufferInfo{ buffer.data(), buffer.size() };
 
-  size_t      bufferInfoAddress = reinterpret_cast<size_t>(&bufferInfo);
-  std::string memAddress = std::to_string(bufferInfoAddress) + ".memory";
+  std::string memAddress = itk::OMEZarrNGFFImageIO::MakeMemoryFileName(bufferInfo);
 
   reader->SetFileName(memAddress);
   ITK_TRY_EXPECT_NO_EXCEPTION(reader->Update());
@@ -83,6 +82,7 @@ doTest(const char * inputFileName, const char * outputFileName)
   writer->SetImageIO(zarrIO);
   ITK_TRY_EXPECT_NO_EXCEPTION(writer->Update());
 
+  // Write zip bitstream to disk
   std::ofstream oFile(outputFileName, std::ios::binary);
   oFile.write(bufferInfo.pointer, bufferInfo.size);
   free(bufferInfo.pointer);