From 461228781d51e01cd0e23d433e5f7974f660288a Mon Sep 17 00:00:00 2001
From: Miles <miles.stoudenmire@gmail.com>
Date: Mon, 6 May 2024 16:28:40 -0400
Subject: [PATCH] Add get! functions and test

---
 src/abstractdatagraph.jl | 14 ++++++++++++++
 test/runtests.jl         | 21 +++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/src/abstractdatagraph.jl b/src/abstractdatagraph.jl
index 41f798c..e03cfef 100644
--- 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
index ab1faae..d58bf89 100644
--- 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