support IO for writetable (#245)

* support IO for writetable * add test for writetable to IO * use io for writetable test --------- Co-authored-by: Helmut Hänsel <helmut.haensel@gmx.com>
felipenoris · Dec 7, 2023 · edba58c · edba58c
1 parent 8e63de8
commit edba58c
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 13 deletions.
diff --git a/src/tables_interface.jl b/src/tables_interface.jl
@@ -27,18 +27,18 @@ end
 
 Write Tables.jl `table` to the specified filename.
 """
-writetable(filename::AbstractString, x; kw...) = writetable(filename, _table_to_arrays(x)...; kw...)
+writetable(filename::Union{AbstractString, IO}, x; kw...) = writetable(filename, _table_to_arrays(x)...; kw...)
 
 """
-    writetable(filename::AbstractString, tables::Vector{Pair{String, T}}; overwrite::Bool=false)
-    writetable(filename::AbstractString, tables::Pair{String, Any}...; overwrite::Bool=false)
+    writetable(filename::Union{AbstractString, IO}, tables::Vector{Pair{String, T}}; overwrite::Bool=false)
+    writetable(filename::Union{AbstractString, IO}, tables::Pair{String, Any}...; overwrite::Bool=false)
 """
-function writetable(filename::AbstractString, tables::Vector{<:Pair}; kw...)
+function writetable(filename::Union{AbstractString, IO}, tables::Vector{<:Pair}; kw...)
     data = [(name, _table_to_arrays(x)...) for (name, x) in tables]
     return writetable(filename, data; kw...)
 end
 
-writetable(filename::AbstractString, tables::Pair{<:String, <:Any}...; kw...) = writetable(filename, collect(tables); kw...)
+writetable(filename::Union{AbstractString, IO}, tables::Pair{<:String, <:Any}...; kw...) = writetable(filename, collect(tables); kw...)
 
 """
     writetable!(sheet::Worksheet, table; anchor_cell::CellRef=CellRef("A1")))

diff --git a/src/write.jl b/src/write.jl
@@ -666,9 +666,9 @@ XLSX.writetable("table.xlsx", columns, colnames)
 
 See also: [`XLSX.writetable!`](@ref).
 """
-function writetable(filename::AbstractString, data, columnnames; overwrite::Bool=false, sheetname::AbstractString="", anchor_cell::Union{String, CellRef}=CellRef("A1"))
+function writetable(filename::Union{AbstractString, IO}, data, columnnames; overwrite::Bool=false, sheetname::AbstractString="", anchor_cell::Union{String, CellRef}=CellRef("A1"))
 
-    if !overwrite
+    if filename isa AbstractString && !overwrite
         @assert !isfile(filename) "$filename already exists."
     end
 
@@ -687,8 +687,8 @@ function writetable(filename::AbstractString, data, columnnames; overwrite::Bool
 end
 
 """
-    writetable(filename::AbstractString; overwrite::Bool=false, kw...)
-    writetable(filename::AbstractString, tables::Vector{Tuple{String, Vector{Any}, Vector{String}}}; overwrite::Bool=false)
+    writetable(filename::Union{AbstractString, IO}; overwrite::Bool=false, kw...)
+    writetable(filename::Union{AbstractString, IO}, tables::Vector{Tuple{String, Vector{Any}, Vector{String}}}; overwrite::Bool=false)
 
 Write multiple tables.
 
@@ -707,9 +707,9 @@ julia> df2 = DataFrames.DataFrame(AA=["aa", "bb"], AB=[10.1, 10.2])
 julia> XLSX.writetable("report.xlsx", "REPORT_A" => df1, "REPORT_B" => df2)
 ```
 """
-function writetable(filename::AbstractString; overwrite::Bool=false, kw...)
+function writetable(filename::Union{AbstractString, IO}; overwrite::Bool=false, kw...)
 
-    if !overwrite
+    if filename isa AbstractString && !overwrite
         @assert !isfile(filename) "$filename already exists."
     end
 
@@ -735,9 +735,9 @@ function writetable(filename::AbstractString; overwrite::Bool=false, kw...)
     nothing
 end
 
-function writetable(filename::AbstractString, tables::Vector{Tuple{String, S, Vector{T}}}; overwrite::Bool=false) where {S<:Vector{U} where U, T<:Union{String, Symbol}}
+function writetable(filename::Union{AbstractString, IO}, tables::Vector{Tuple{String, S, Vector{T}}}; overwrite::Bool=false) where {S<:Vector{U} where U, T<:Union{String, Symbol}}
 
-    if !overwrite
+    if filename isa AbstractString && !overwrite
         @assert !isfile(filename) "$filename already exists."
     end
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1314,6 +1314,17 @@ end
         check_test_data(data, report_2_data)
     end
 
+    @testset "writetable to IO" begin
+        dt = XLSX.DataTable(Any[Any[1, 2, 3], Any[4, 5, 6]], [:a, :b])
+        io = IOBuffer()
+        XLSX.writetable(io, "Test" => dt)
+        seek(io, 0)
+        dt_read = XLSX.readtable(io, "Test")
+        @test dt_read.data == dt.data
+        @test dt_read.column_labels == dt.column_labels
+        @test dt_read.column_label_index == dt.column_label_index
+    end
+
     # delete files created by this testset
     delete_files = ["output_table.xlsx", "output_tables.xlsx"]
     for f in delete_files