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/MOF/MOF.jl

@@ -90,8 +90,9 @@ Validate that the MOF file `filename` conforms to the MOF JSON schema. Returns
 `nothing` if the file is valid, otherwise throws an error describing why the
 file is not valid.
 """
-function validate(filename::String)
-    MathOptFormat.gzip_open(filename, "r") do io
+function validate(filename::String; compression::MathOptFormat.AbstractCompressionScheme=MathOptFormat.AutomaticCompression())
+    compression = MathOptFormat._automatic_compression(filename, compression)
+    MathOptFormat._compressed_open(filename, "r", compression) do io
         validate(io)
     end
     return

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,51 @@ 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
 
 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=AutomaticCompression())
+    compression = _automatic_compression(filename, compression)
+    _compressed_open(filename, "w", 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=AutomaticCompression())
+    compression = _automatic_compression(filename, compression)
+    _compressed_open(filename, "r", compression) do io
         MOI.read_from_file(model, io)
     end
 end
@@ -173,19 +191,9 @@ end
 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
-    MOI.read_from_file(model, filename)
+function read_from_file(filename::String; compression::AbstractCompressionScheme=AutomaticCompression())
+    model = _filename_to_model(filename)
+    MOI.read_from_file(model, filename, compression=compression)
     return model
 end
 

src/compression.jl

@@ -0,0 +1,80 @@
+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 AutomaticCompression <: 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 _compressed_open(
+    f::Function, filename::String, mode::String, ::NoCompression
+)
+    return Base.open(f, filename, mode)
+end
+
+struct Gzip <: AbstractCompressionScheme end
+function _compressed_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 _compressed_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 _compressed_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 _automatic_compression(filename::String, compression::AbstractCompressionScheme)
+    if compression == AutomaticCompression()
+        return _filename_to_compression(filename)
+    end
+    return compression
+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,15 +27,35 @@ const MOIU = MOI.Utilities
         end
     end
 
-    @testset "Calling gzip_open" begin
-        @test_throws ArgumentError MathOptFormat.gzip_open((x) -> nothing,
-                                                           "dummy.gz", "a")
-        @test_throws ArgumentError MathOptFormat.gzip_open((x) -> nothing,
-                                                           "dummy.gz", "r+")
-        @test_throws ArgumentError MathOptFormat.gzip_open((x) -> nothing,
-                                                           "dummy.gz", "w+")
-        @test_throws ArgumentError MathOptFormat.gzip_open((x) -> nothing,
-                                                           "dummy.gz", "a+")
-
+    @testset "Calling MOF._compressed_open" begin
+        for cs in [MathOptFormat.Bzip2(), MathOptFormat.Gzip()] # MathOptFormat.Xz()
+            @test_throws ArgumentError MathOptFormat._compressed_open((x) -> nothing,
+                                                          "dummy.gz", "a", cs)
+            @test_throws ArgumentError MathOptFormat._compressed_open((x) -> nothing,
+                                                          "dummy.gz", "r+", cs)
+            @test_throws ArgumentError MathOptFormat._compressed_open((x) -> nothing,
+                                                          "dummy.gz", "w+", cs)
+            @test_throws ArgumentError MathOptFormat._compressed_open((x) -> nothing,
+                                                          "dummy.gz", "a+", cs)
+        end
+    end
+
+    @testset "Provided compression schemes" begin
+        file_to_read = joinpath(@__DIR__, "MPS", "free_integer.mps")
+        m = MathOptFormat.read_from_file(file_to_read)
+
+        @testset "Automatic detection from extension" begin
+            MOI.write_to_file(m, file_to_read * ".garbage")
+            for ext in ["", ".bz2", ".gz"] # ".xz"
+                MOI.write_to_file(m, file_to_read * ext)
+                MathOptFormat.read_from_file(file_to_read * ext)
+            end
+
+            # Clean up
+            sleep(1.0)  # Allow time for unlink to happen.
+            for ext in ["", ".garbage", ".bz2", ".gz"] # ".xz"
+                rm(file_to_read * ext, force = true)
+            end
+        end
     end
 end