From d0d3098fcae38a1de2b1a71b8eaccfbe7940cf7d Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Wed, 17 Apr 2024 09:06:45 -0400
Subject: [PATCH] Allow specifying vertex type

---
 Project.toml     |  2 +-
 src/datagraph.jl | 17 +++++++++++------
 test/runtests.jl | 42 ++++++++++++++++++++++++++++++++++++++----
 3 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/Project.toml b/Project.toml
index 3864a4c..b4a46af 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "DataGraphs"
 uuid = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org> and contributors"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
diff --git a/src/datagraph.jl b/src/datagraph.jl
index 718bf08..b6d1784 100644
--- a/src/datagraph.jl
+++ b/src/datagraph.jl
@@ -1,7 +1,7 @@
 using Dictionaries: Dictionary
 using Graphs: Graphs, edgetype
 using Graphs.SimpleGraphs: SimpleGraph
-using NamedGraphs.GraphsExtensions: directed_graph, vertextype
+using NamedGraphs.GraphsExtensions: convert_vertextype, directed_graph, vertextype
 
 # TODO: define VertexDataGraph, a graph with only data on the
 # vertices, and EdgeDataGraph, a graph with only data on the edges.
@@ -46,16 +46,21 @@ function Base.copy(graph::DataGraph)
   )
 end
 
-function DataGraph(
+function DataGraph{V}(
   underlying_graph::AbstractGraph; vertex_data_eltype::Type=Any, edge_data_eltype::Type=Any
-)
+) where {V}
+  converted_underlying_graph = convert_vertextype(V, underlying_graph)
   return _DataGraph(
-    underlying_graph,
-    Dictionary{vertextype(underlying_graph),vertex_data_eltype}(),
-    Dictionary{edgetype(underlying_graph),edge_data_eltype}(),
+    converted_underlying_graph,
+    Dictionary{vertextype(converted_underlying_graph),vertex_data_eltype}(),
+    Dictionary{edgetype(converted_underlying_graph),edge_data_eltype}(),
   )
 end
 
+function DataGraph(underlying_graph::AbstractGraph; kwargs...)
+  return DataGraph{vertextype(underlying_graph)}(underlying_graph; kwargs...)
+end
+
 function DataGraph{V,VD,ED,G,E}(underlying_graph::AbstractGraph) where {V,VD,ED,G,E}
   @assert edgetype(underlying_graph) === E
   return _DataGraph(convert(G, underlying_graph), Dictionary{V,VD}(), Dictionary{E,ED}())
diff --git a/test/runtests.jl b/test/runtests.jl
index bb1f738..da410d1 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,6 +1,13 @@
 @eval module $(gensym())
-using DataGraphs: DataGraphs, DataGraph, is_arranged
-using Dictionaries: Indices, dictionary
+using DataGraphs:
+  DataGraphs,
+  DataGraph,
+  edge_data,
+  edge_data_eltype,
+  is_arranged,
+  vertex_data,
+  vertex_data_eltype
+using Dictionaries: Dictionary, Indices, dictionary
 using Graphs:
   add_edge!,
   bfs_tree,
@@ -10,6 +17,7 @@ using Graphs:
   dijkstra_shortest_paths,
   dst,
   edges,
+  edgetype,
   grid,
   has_edge,
   has_vertex,
@@ -22,8 +30,8 @@ using Graphs:
   vertices
 using Graphs.SimpleGraphs: SimpleDiGraph, SimpleEdge, SimpleGraph
 using GraphsFlows: GraphsFlows
-using NamedGraphs: NamedDiGraph, NamedGraph
-using NamedGraphs.GraphsExtensions: ⊔, rename_vertices
+using NamedGraphs: NamedDiGraph, NamedEdge, NamedGraph
+using NamedGraphs.GraphsExtensions: ⊔, rename_vertices, vertextype
 using NamedGraphs.NamedGraphGenerators: named_grid, named_path_graph
 using Test: @test, @test_broken, @testset
 
@@ -175,6 +183,32 @@ using DataGraphs: is_arranged
     @test dg["Y" => 2] == "edge_Y2"
   end
 
+  @testset "Constructors specifying vertex type" begin
+    dg = DataGraph{Float64}(
+      named_path_graph(4); vertex_data_eltype=String, edge_data_eltype=Symbol
+    )
+    @test nv(dg) == 4
+    @test ne(dg) == 3
+    @test edgetype(dg) === NamedEdge{Float64}
+    @test vertextype(dg) === Float64
+    @test vertex_data_eltype(dg) === String
+    @test edge_data_eltype(dg) === Symbol
+    @test issetequal(vertices(dg), Float64.(1:4))
+    @test vertices(dg) isa Indices{Float64}
+    @test eltype(vertices(dg)) === Float64
+    @test has_edge(dg, 1.0 => 2.0)
+    @test has_edge(dg, 2.0 => 3.0)
+    @test has_edge(dg, 3.0 => 4.0)
+    @test vertex_data(dg) == Dictionary{Float64,String}()
+    @test vertex_data(dg) isa Dictionary{Float64,String}
+    @test keytype(vertex_data(dg)) === Float64
+    @test eltype(vertex_data(dg)) === String
+    @test edge_data(dg) == Dictionary{NamedEdge{Float64},Symbol}()
+    @test edge_data(dg) isa Dictionary{NamedEdge{Float64},Symbol}
+    @test keytype(edge_data(dg)) === NamedEdge{Float64}
+    @test eltype(edge_data(dg)) === Symbol
+  end
+
   @testset "Disjoint unions" begin
     g = DataGraph(named_grid((2, 2)); vertex_data_eltype=String, edge_data_eltype=String)