Add support for LZMA compression of G3 files

Author: arahlin

core/CMakeLists.txt

@@ -52,6 +52,16 @@ if(BZIP2_FOUND)
 	target_compile_definitions(core PRIVATE -DBZIP2_FOUND)
 endif()
 
+if(NOT DEFINED WITH_LZMA)
+	set(WITH_LZMA TRUE CACHE BOOL "Enable LZMA file compression")
+endif()
+if(WITH_LZMA)
+	find_package(LibLZMA)
+endif()
+if(LIBLZMA_FOUND)
+	target_compile_definitions(core PRIVATE -DLZMA_FOUND)
+endif()
+
 link_python_dir()
 
 add_spt3g_program(bin/spt3g-dump)
@@ -80,6 +90,9 @@ add_spt3g_test(compressedfileio)
 if(BZIP2_FOUND)
 	add_spt3g_test(bz2fileio)
 endif()
+if(LIBLZMA_FOUND)
+	add_spt3g_test(xzfileio)
+endif()
 add_spt3g_test(portability)
 add_spt3g_test(vecint)
 add_spt3g_test(ts_bufferprotocol)

core/src/dataio.cxx

@@ -5,6 +5,9 @@
 #ifdef BZIP2_FOUND
 #include <boost/iostreams/filter/bzip2.hpp>
 #endif
+#ifdef LZMA_FOUND
+#include <boost/iostreams/filter/lzma.hpp>
+#endif
 #include <boost/iostreams/device/file.hpp>
 #include <boost/iostreams/device/file_descriptor.hpp>
 #include <boost/iostreams/device/array.hpp>
@@ -33,6 +36,13 @@ g3_istream_from_path(g3_istream &stream, const std::string &path,
 		log_fatal("Boost not compiled with bzip2 support.");
 #endif
 	}
+	if (path.size() > 3 && !path.compare(path.size() - 3, 3, ".xz")) {
+#ifdef LZMA_FOUND
+		stream.push(boost::iostreams::lzma_decompressor());
+#else
+		log_fatal("Boost not compiled with LZMA support.");
+#endif
+	}
 
 	int fd = -1;
 
@@ -189,6 +199,13 @@ g3_ostream_to_path(g3_ostream &stream, const std::string &path,
 		log_fatal("Boost not compiled with bzip2 support.");
 #endif
 	}
+	if (path.size() > 3 && !path.compare(path.size() - 3, 3, ".xz") && !append) {
+#ifdef LZMA_FOUND
+		stream.push(boost::iostreams::lzma_compressor());
+#else
+		log_fatal("Boost not compiled with LZMA support.");
+#endif
+	}
 
 	if (counter)
 		stream.push(boost::iostreams::counter64());

core/tests/xzfileio.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+
+from spt3g import core
+import time
+
+# File to disk
+pipe = core.G3Pipeline()
+pipe.Add(core.G3InfiniteSource, type=core.G3FrameType.Timepoint, n=10)
+n = 0
+
+
+def addinfo(fr):
+    global n
+    if fr.type != core.G3FrameType.Timepoint:
+        return
+    fr["time"] = core.G3Time(int(time.time() * core.G3Units.s))
+    fr["count"] = n
+    n += 1
+
+
+pipe.Add(addinfo)
+pipe.Add(core.Dump)
+pipe.Add(core.G3Writer, filename="test.g3.xz")
+pipe.Run()
+
+# And back from disk
+print("Reading")
+pipe = core.G3Pipeline()
+pipe.Add(core.G3Reader, filename="test.g3.xz")
+pipe.Add(core.Dump)
+n = 0
+
+
+def checkinfo(fr):
+    global n
+    if fr.type != core.G3FrameType.Timepoint:
+        return
+    if "time" not in fr:
+        raise KeyError("time")
+    if fr["count"] != n:
+        raise ValueError("Out of order frame")
+    n += 1
+
+
+pipe.Add(checkinfo)
+pipe.Run()
+
+if n != 10:
+    raise ValueError("Wrong number of frames (%d should be %d)" % (n, 10))