Add more compression formats

Project.toml

@@ -4,6 +4,8 @@ authors = ["Oscar Dowson <[email protected]"]
 version = "0.2.2"
 
 [deps]
+CodecBzip2 = "523fee87-0ab8-5b00-afb7-3ecf72e48cfd"
+CodecXz = "ba30903b-d9e8-5048-a5ec-d1f5b0d4b47b"
 CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"

src/MathOptFormat.jl

@@ -3,8 +3,11 @@ module MathOptFormat
 import MathOptInterface
 const MOI = MathOptInterface
 
+import CodecBzip2
+import CodecXz
 import CodecZlib
 
+include("compression.jl")
 include("CBF/CBF.jl")
 include("LP/LP.jl")
 include("MOF/MOF.jl")
@@ -133,36 +136,56 @@ function create_unique_variable_names(
     end
 end
 
-function gzip_open(f::Function, filename::String, mode::String)
-    if endswith(filename, ".gz")
-        if mode == "r"
-            open(CodecZlib.GzipDecompressorStream, filename, mode) do io
-                f(io)
-            end
-        elseif mode == "w"
-            open(CodecZlib.GzipCompressorStream, filename, mode) do io
-                f(io)
+"""
+List of accepted export formats. `AUTOMATIC_FILE_FORMAT` corresponds to
+a detection from the file name, only based on the extension (regardless of
+compression format).
+"""
+@enum(FileFormat, FORMAT_CBF, FORMAT_LP, FORMAT_MOF, FORMAT_MPS, AUTOMATIC_FILE_FORMAT)
+
+const _file_formats = Dict{FileFormat, Tuple{String, Any}}(
+    # ENUMERATED VALUE => extension, model type
+    FORMAT_CBF => (".cbf", CBF.Model),
+    FORMAT_LP => (".lp", LP.Model),
+    FORMAT_MOF => (".mof.json", MOF.Model),
+    FORMAT_MPS => (".mps", MPS.Model)
+)
+
+function _filename_to_format(filename::String)
+    for compr_ext in ["", ".bz2", ".gz", ".xz"]
+        for (type, format) in _file_formats
+            if endswith(filename, "$(format[1])$(compr_ext)")
+                return type
             end
-        else
-            throw(ArgumentError("Mode must be \"r\" or \"w\""))
         end
-    else
-        return open(f, filename, mode)
     end
+
+    error("File type of $(filename) not recognized by MathOptFormat.jl.")
+end
+
+function _filename_to_model(filename::String)
+    return _file_formats[_filename_to_format(filename)][2]()
+end
+
+function gzip_open(f::Function, filename::String, mode::String; compression::AbstractCompressionScheme=AutomaticCompressionDetection())
+    if compression == AutomaticCompressionDetection()
+        compression = _filename_to_compression(filename)
+    end
+    return open(f, filename, mode, compression)
 end
 
 const MATH_OPT_FORMATS = Union{
     CBF.InnerModel, LP.InnerModel, MOF.Model, MPS.InnerModel
 }
 
-function MOI.write_to_file(model::MATH_OPT_FORMATS, filename::String)
-    gzip_open(filename, "w") do io
+function MOI.write_to_file(model::MATH_OPT_FORMATS, filename::String; compression::AbstractCompressionScheme=AutomaticCompressionDetection())
+    gzip_open(filename, "w", compression=compression) do io
         MOI.write_to_file(model, io)
     end
 end
 
-function MOI.read_from_file(model::MATH_OPT_FORMATS, filename::String)
-    gzip_open(filename, "r") do io
+function MOI.read_from_file(model::MATH_OPT_FORMATS, filename::String; compression::AbstractCompressionScheme=AutomaticCompressionDetection())
+    gzip_open(filename, "r", compression=compression) do io
         MOI.read_from_file(model, io)
     end
 end
@@ -174,17 +197,7 @@ Create a MOI model by reading `filename`. Type of the returned model depends on
 the extension of `filename`.
 """
 function read_from_file(filename::String)
-    model = if endswith(filename, ".mof.json.gz") || endswith(filename, ".mof.json")
-        MOF.Model()
-    elseif endswith(filename, ".cbf.gz") || endswith(filename, ".cbf")
-        CBF.Model()
-    elseif endswith(filename, ".mps.gz") || endswith(filename, ".mps")
-        MPS.Model()
-    elseif endswith(filename, ".lp.gz") || endswith(filename, ".lp")
-        LP.Model()
-    else
-        error("File-type of $(filename) not supported by MathOptFormat.jl.")
-    end
+    model = _filename_to_model(filename)
     MOI.read_from_file(model, filename)
     return model
 end

src/compression.jl

@@ -0,0 +1,73 @@
+function error_mode(mode::String)
+    throw(ArgumentError("For dealing with compressed data, mode must be \"r\" or \"w\"; $mode given"))
+end
+
+"""
+    abstract type AbstractCompressionScheme end
+
+Base type to implement a new compression scheme for MathOptFormat. To do so,
+create a concrete subtype (e.g., named after the compression scheme) and
+implement `open(f::Function, filename::String, mode::String, ::YourScheme)`. 
+"""
+abstract type AbstractCompressionScheme end
+
+struct AutomaticCompressionDetection <: AbstractCompressionScheme end
+# No open() implementation, this would not make sense (flag to indicate that _filename_to_compression should be called).
+
+struct NoCompression <: AbstractCompressionScheme end
+function open(
+    f::Function, filename::String, mode::String, ::NoCompression
+)
+    return Base.open(f, filename, mode)
+end
+
+struct Gzip <: AbstractCompressionScheme end
+function open(
+    f::Function, filename::String, mode::String, ::Gzip
+)
+    return if mode == "w"
+        Base.open(f, CodecZlib.GzipCompressorStream, filename, mode)
+    elseif mode == "r"
+        Base.open(f, CodecZlib.GzipDecompressorStream, filename, mode)
+    else
+        error_mode(mode)
+    end
+end
+
+struct Bzip2 <: AbstractCompressionScheme end
+function open(
+    f::Function, filename::String, mode::String, ::Bzip2
+)
+    if mode == "w"
+        Base.open(f, CodecBzip2.Bzip2CompressorStream, filename, mode)
+    elseif mode == "r"
+        Base.open(f, CodecBzip2.Bzip2DecompressorStream, filename, mode)
+    else
+        error_mode(mode)
+    end
+end
+
+struct Xz <: AbstractCompressionScheme end
+function open(
+    f::Function, filename::String, mode::String, ::Xz
+)
+    return if mode == "w"
+        Base.open(f, CodecXz.XzDecompressorStream, filename, mode)
+    elseif mode == "r"
+        Base.open(f, CodecXz.XzCompressorStream, filename, mode)
+    else
+        error_mode(mode)
+    end
+end
+
+function _filename_to_compression(filename::String)
+    if endswith(filename, ".bz2")
+        return Bzip2()
+    elseif endswith(filename, ".gz")
+        return Gzip()
+    elseif endswith(filename, ".xz")
+        return Xz()
+    else
+        return NoCompression()
+    end
+end

test/runtests.jl

@@ -27,6 +27,19 @@ const MOIU = MOI.Utilities
         end
     end
 
+    @testset "Calling MOF.open" begin
+        for cs in [MathOptFormat.Bzip2(), MathOptFormat.Gzip(), MathOptFormat.Xz()]
+            @test_throws ArgumentError MathOptFormat.open((x) -> nothing,
+                                                          "dummy.gz", "a", cs)
+            @test_throws ArgumentError MathOptFormat.open((x) -> nothing,
+                                                          "dummy.gz", "r+", cs)
+            @test_throws ArgumentError MathOptFormat.open((x) -> nothing,
+                                                          "dummy.gz", "w+", cs)
+            @test_throws ArgumentError MathOptFormat.open((x) -> nothing,
+                                                          "dummy.gz", "a+", cs)
+        end
+    end
+
     @testset "Calling gzip_open" begin
         @test_throws ArgumentError MathOptFormat.gzip_open((x) -> nothing,
                                                            "dummy.gz", "a")
@@ -36,6 +49,5 @@ const MOIU = MOI.Utilities
                                                            "dummy.gz", "w+")
         @test_throws ArgumentError MathOptFormat.gzip_open((x) -> nothing,
                                                            "dummy.gz", "a+")
-
     end
 end