Typecast GNNGraph.num_nodes to Int (#109)

JuliaGraphs · Jan 20, 2022 · df5b508 · df5b508
1 parent ce0396b
commit df5b508
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git a/src/GNNGraphs/convert.jl b/src/GNNGraphs/convert.jl
@@ -2,7 +2,7 @@
 
 function to_coo(coo::COO_T; dir=:out, num_nodes=nothing, weighted=true)
     s, t, val = coo   
-    num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
+    num_nodes::Int = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
     @assert isnothing(val) || length(val) == length(s)
     @assert length(s) == length(t)
     if !isempty(s)
@@ -114,7 +114,7 @@ function to_dense(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing, weighted=t
     # `dir` will be ignored since the input `coo` is always in source -> target format.
     # The output will always be a adjmat in :out format (e.g. A[i,j] denotes from i to j)
     s, t, val = coo
-    n = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes
+    n::Int = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes
     val = isnothing(val) ? eltype(s)(1) : val
     T = T === nothing ? eltype(val) : T
     if !weighted
@@ -164,9 +164,9 @@ function to_sparse(coo::COO_T, T=nothing; dir=:out, num_nodes=nothing, weighted=
         eweight = fill!(similar(s, T), 1)
     end
 
-    num_nodes = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
+    num_nodes::Int = isnothing(num_nodes) ? max(maximum(s), maximum(t)) : num_nodes 
     A = sparse(s, t, eweight, num_nodes, num_nodes)
-    num_edges = nnz(A)
+    num_edges::Int = nnz(A)
     if eltype(A) != T
         A = T.(A)
     end

diff --git a/src/GNNGraphs/gnngraph.jl b/src/GNNGraphs/gnngraph.jl
@@ -168,7 +168,8 @@ function GNNGraph(g::AbstractGraph; kws...)
         # add reverse edges since GNNGraph is directed
         s, t = [s; t], [t; s]    
     end
-    GNNGraph((s, t); num_nodes=Graphs.nv(g), kws...)
+    num_nodes::Int = Graphs.nv(g)
+    GNNGraph((s, t); num_nodes=num_nodes, kws...)
 end
 
 

diff --git a/test/GNNGraphs/gnngraph.jl b/test/GNNGraphs/gnngraph.jl
@@ -177,7 +177,12 @@
         for e in Graphs.edges(lg)
             i, j = src(e), dst(e)
             @test has_edge(g, i, j)
-            @test has_edge(g, j, i)            
+            @test has_edge(g, j, i)
+        end
+
+        @testset "SimpleGraph{Int32}" begin
+            g = GNNGraph(SimpleGraph{Int32}(6), graph_type=GRAPH_T)
+            @test g.num_nodes == 6
         end
     end