diff --git a/src/LoadSave/circuit_savers.jl b/src/LoadSave/circuit_savers.jl index 75d6acfb..d8d74f14 100644 --- a/src/LoadSave/circuit_savers.jl +++ b/src/LoadSave/circuit_savers.jl @@ -4,11 +4,15 @@ export save_as_dot, save_circuit, save_as_sdd, save_as_tex, save_as_dot2tex, sav # Save lines ##################### -function save_lines(name::String, lines::CircuitFormatLines) - open(name, "w") do f - for line in lines - println(f, line) +function save_lines(file::Union{String, IO}, lines::CircuitFormatLines) + if file isa String + open(file, "w") do f + for line in lines + println(f, line) + end end + else + for line in lines println(file, line) end end end diff --git a/src/LoadSave/vtree_loaders.jl b/src/LoadSave/vtree_loaders.jl index ff084d66..a9189c82 100644 --- a/src/LoadSave/vtree_loaders.jl +++ b/src/LoadSave/vtree_loaders.jl @@ -5,7 +5,7 @@ export load_vtree, zoo_vtree, zoo_vtree_file Load a vtree file from file. Currently only supports ".vtree" format. """ -function load_vtree(file::String, ::Type{V}=PlainVtree)::V where V<:Vtree +function load_vtree(file::Union{String, IO}, ::Type{V}=PlainVtree)::V where V<:Vtree return compile_vtree_format_lines(parse_vtree_file(file), V) end @@ -41,22 +41,22 @@ function parse_vtree_header_line(ln::String) VtreeHeaderLine() end -function parse_vtree_file(file::String)::VtreeFormatLines +function parse_vtree_file(file::Union{String, IO})::VtreeFormatLines q = Vector{VtreeFormatLine}() - open(file) do file - for ln in eachline(file) - if ln[1] == 'c' - push!(q, parse_vtree_comment_line(ln)) - elseif ln[1] == 'L' - push!(q, parse_vtree_leaf_line(ln)) - elseif ln[1] == 'I' - push!(q, parse_vtree_inner_line(ln)) - elseif startswith(ln, "vtree") - push!(q, parse_vtree_header_line(ln)) - else - error("Don't know how to parse vtree file format line $ln") - end + if file isa String file = open(file) end + for ln in eachline(file) + if ln[1] == 'c' + push!(q, parse_vtree_comment_line(ln)) + elseif ln[1] == 'L' + push!(q, parse_vtree_leaf_line(ln)) + elseif ln[1] == 'I' + push!(q, parse_vtree_inner_line(ln)) + elseif startswith(ln, "vtree") + push!(q, parse_vtree_header_line(ln)) + else + error("Don't know how to parse vtree file format line $ln") end end + close(file) q -end \ No newline at end of file +end diff --git a/src/LoadSave/vtree_savers.jl b/src/LoadSave/vtree_savers.jl index 17aefa2e..2281452a 100644 --- a/src/LoadSave/vtree_savers.jl +++ b/src/LoadSave/vtree_savers.jl @@ -50,16 +50,18 @@ Supported formats: * ".vtree" for Vtree files * ".dot" for dot files """ -function save_vtree(file::AbstractString, vtree::PlainVtree) +function save_vtree(file::Union{AbstractString, IO}, vtree::PlainVtree) # 1. decide file type and open file - if endswith(file,".vtree") - f = VtreeConfigFile(file) - elseif endswith(file, ".dot") - f = VtreeDotFile(file) - else - throw("Unsupported file type. For example, for vtree format file name should end with .vtree") - end + if file isa AbstractString + if endswith(file,".vtree") + f = VtreeConfigFile(file) + elseif endswith(file, ".dot") + f = VtreeDotFile(file) + else + throw("Unsupported file type. For example, for vtree format file name should end with .vtree") + end + else f = (file = file,) end #2. map from PlainVtree to index for output index_cache = Dict{PlainVtree, UInt32}() @@ -70,7 +72,7 @@ function save_vtree(file::AbstractString, vtree::PlainVtree) end #3. saving methods for header, nodes, tailer - function save_vtree_header(vtree::PlainVtree, f::VtreeConfigFile) + function save_vtree_header(vtree::PlainVtree, f::Union{VtreeConfigFile, NamedTuple}) write(f.file, VTREE_FORMAT) write(f.file, "vtree $(num_nodes(vtree))\n") end @@ -83,7 +85,7 @@ function save_vtree(file::AbstractString, vtree::PlainVtree) node_index = node2index(n) node_variable = n.var - if f isa VtreeConfigFile + if f isa VtreeConfigFile || f isa NamedTuple write(f.file, "L $node_index $node_variable\n") elseif f isa VtreeDotFile write(f.file, "$node_index [label=$(node_variable), shape=\"plaintext\"]\n") @@ -97,7 +99,7 @@ function save_vtree(file::AbstractString, vtree::PlainVtree) left = node2index(n.left) right = node2index(n.right) - if f isa VtreeConfigFile + if f isa VtreeConfigFile || f isa NamedTuple write(f.file, "I $node_index $left $right\n") elseif f isa VtreeDotFile write(f.file, "$node_index -- $right\n") @@ -107,7 +109,7 @@ function save_vtree(file::AbstractString, vtree::PlainVtree) end end - function save_vtree_footer(f::VtreeConfigFile) + function save_vtree_footer(f::Union{VtreeConfigFile, NamedTuple}) end #do nothing function save_vtree_footer(f::VtreeDotFile) diff --git a/test/LoadSave/vtree_parser_test.jl b/test/LoadSave/vtree_parser_test.jl index f8bdb5d9..f67e92f7 100644 --- a/test/LoadSave/vtree_parser_test.jl +++ b/test/LoadSave/vtree_parser_test.jl @@ -33,16 +33,25 @@ temp_path = "$tmp/little_4var_temp.vtree" save_vtree(temp_path, vtree) - # load it from file, and then run the same tests - vtree2 = load_vtree(temp_path) - test_vtree(vtree2) - @test vtree == vtree2 # we can test equality of plain vtrees! + # Save vtree from existing file stream + temp_stream = open("$tmp/little_4var_temp_stream.vtree", "w") + save_vtree(temp_stream, vtree) + close(temp_stream) + + # load from file, and then run the same tests + temp_stream = open("$tmp/little_4var_temp_stream.vtree", "r") + for f in [temp_path, temp_stream] + vtree2 = load_vtree(f) + test_vtree(vtree2) + @test vtree == vtree2 # we can test equality of plain vtrees! + end + close(temp_stream) # Save dot file dot_path = "$tmp/little_4var_temp.dot" save_vtree(dot_path, vtree) - # Save unsupported format + # Save unsupported format dot_path = "$tmp/little_4var_temp.bad_extension" @test_throws String save_vtree(dot_path, vtree) end