Add get! functions and test (#33)

ITensor · May 7, 2024 · 3e5b471 · 3e5b471
1 parent 0f8075d
commit 3e5b471
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 0 deletions.
diff --git a/src/abstractdatagraph.jl b/src/abstractdatagraph.jl
@@ -306,6 +306,10 @@ function Base.get(graph::AbstractDataGraph, vertex, default)
   return get(vertex_data(graph), vertex, default)
 end
 
+function Base.get!(graph::AbstractDataGraph, vertex, default)
+  return get!(vertex_data(graph), vertex, default)
+end
+
 function Base.getindex(graph::AbstractDataGraph, edge::AbstractEdge)
   is_edge_arranged_ = is_edge_arranged(graph, edge)
   data = edge_data(graph)[arrange(is_edge_arranged_, edge)]
@@ -327,6 +331,16 @@ function Base.get(graph::AbstractDataGraph, edge::Pair, default)
   return get(graph, edgetype(graph)(edge), default)
 end
 
+function Base.get!(graph::AbstractDataGraph, edge::AbstractEdge, default)
+  is_edge_arranged_ = is_edge_arranged(graph, edge)
+  data = get!(edge_data(graph), arrange(is_edge_arranged_, edge), default)
+  return reverse_data_direction(is_edge_arranged_, graph, data)
+end
+
+function Base.get!(graph::AbstractDataGraph, edge::Pair, default)
+  return get!(graph, edgetype(graph)(edge), default)
+end
+
 # Support syntax `g[1, 2] = g[(1, 2)]`
 function Base.getindex(graph::AbstractDataGraph, i1, i2, i...)
   return graph[(i1, i2, i...)]

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -198,6 +198,27 @@ using DataGraphs: is_arranged
     @test dg["Y" => 2] == "edge_Y2"
   end
 
+  @testset "get and get! functions" begin
+    g = grid((4,))
+    dg = DataGraph(g; vertex_data_eltype=String, edge_data_eltype=Symbol)
+
+    # Test for vertices
+    @test get(dg, 1, "default") == "default"
+    @test !isassigned(dg, 1)
+
+    @test get!(dg, 2, "default") == "default"
+    @test isassigned(dg, 2)
+    @test dg[2] == "default"
+
+    # Test for edges
+    @test get(dg, 1 => 2, :default) == :default
+    @test !isassigned(dg, 1 => 2)
+
+    @test get!(dg, 1 => 2, :default) == :default
+    @test isassigned(dg, 1 => 2)
+    @test dg[1 => 2] == :default
+  end
+
   @testset "Constructors specifying vertex type" begin
     dg = DataGraph{Float64}(
       named_path_graph(4); vertex_data_eltype=String, edge_data_eltype=Symbol