Improve hetero show function (#237)

* Solve indentation problem

* Add test show

Author: aurorarossi

src/GNNGraphs/gnnheterograph.jl

@@ -5,28 +5,28 @@ const NDict{T} = Dict{Symbol, T}
 
 """
     GNNHeteroGraph(data; ndata, edata, gdata, num_nodes, graph_indicator, dir])
-    
+
 A type representing a heterogeneous graph structure.
 it is similar [`GNNGraph`](@ref) but node and edges are of different types.
 
 # Arguments
 
-- `data`: A dictionary or an iterable object that maps (source_type, edge_type, target_type) 
+- `data`: A dictionary or an iterable object that maps (source_type, edge_type, target_type)
     triples to (source, target) index vectors.
 - `num_nodes`: The number of nodes for each type. If not specified, inferred from `g`. Default `nothing`.
-- `graph_indicator`: For batched graphs, a dictionary of vectors containing the graph assignment of each node. Default `nothing`.  
-- `ndata`: Node features. A dictionary of arrays or named tuple of arrays. 
+- `graph_indicator`: For batched graphs, a dictionary of vectors containing the graph assignment of each node. Default `nothing`.
+- `ndata`: Node features. A dictionary of arrays or named tuple of arrays.
            The size of the last dimension of each array must be given by `g.num_nodes`.
 - `edata`: Edge features. A dictionary of arrays or named tuple of arrays.
            The size of the last dimension of each array must be given by `g.num_edges`.
-- `gdata`: Graph features. An array or named tuple of arrays whose last dimension has size `num_graphs`. 
+- `gdata`: Graph features. An array or named tuple of arrays whose last dimension has size `num_graphs`.
 
 
 !!! warning
     `GNNHeteroGraph` is still experimental and not fully supported.
     The interface could be subject to change in the future.
 
-# Examples 
+# Examples
 
 ```julia
 julia> using Flux, GraphNeuralNetworks
@@ -43,7 +43,7 @@ julia> eindex = ((:A, :rel1, :B) => edges1, (:B, :rel2, :A) => edges2);
 
 julia> hg = GNNHeteroGraph(eindex; num_nodes)
 GNNHeteroGraph:
-  num_nodes: (:A => 10, :B => 20)         
+  num_nodes: (:A => 10, :B => 20)
   num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)
 
 julia> hg.num_edges
@@ -57,7 +57,7 @@ julia> ndata = Dict(:A => (x = rand(2, num_nodes[:A]), y = rand(3, num_nodes[:A]
 
 julia> hg = GNNHeteroGraph(eindex; num_nodes, ndata)
 GNNHeteroGraph:
-    num_nodes: (:A => 10, :B => 20)         
+    num_nodes: (:A => 10, :B => 20)
     num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)
     ndata:
     :A  =>  (x = 2×10 Matrix{Float64}, y = 3×10 Matrix{Float64})
@@ -88,15 +88,15 @@ end
 
 @functor GNNHeteroGraph
 
-function GNNHeteroGraph(data::EDict; 
-                        num_nodes = nothing,
-                        graph_indicator = nothing, 
-                        graph_type = :coo,
-                        dir = :out,
-                        ndata = NDict{NamedTuple}(), 
-                        edata = EDict{NamedTuple}(), 
-                        gdata = (;))
-                        
+function GNNHeteroGraph(data::EDict;
+    num_nodes = nothing,
+    graph_indicator = nothing,
+    graph_type = :coo,
+    dir = :out,
+    ndata = NDict{NamedTuple}(),
+    edata = EDict{NamedTuple}(),
+    gdata = (;))
+
 
     @assert graph_type ∈ [:coo, :dense, :sparse] "Invalid graph_type $graph_type requested"
     @assert dir ∈ [:in, :out]
@@ -116,18 +116,18 @@ function GNNHeteroGraph(data::EDict;
     elseif graph_type == :sparse
         graph, num_nodes, num_edges = to_sparse(data; num_nodes, dir)
     end
-    
+
     num_graphs = !isnothing(graph_indicator) ? maximum([maximum(gi) for gi in values(graph_indicator)]) : 1
-    
+
     ndata = normalize_heterographdata(ndata, default_name=:x, n=num_nodes)
     edata = normalize_heterographdata(edata, default_name=:e, n=num_edges, duplicate_if_needed=true)
     gdata = normalize_graphdata(gdata, default_name=:u, n=num_graphs)
-    
-    return GNNHeteroGraph(graph, 
-                num_nodes, num_edges, num_graphs, 
-                graph_indicator,
-                ndata, edata, gdata,
-                ntypes, etypes)
+
+    return GNNHeteroGraph(graph,
+        num_nodes, num_edges, num_graphs,
+        graph_indicator,
+        ndata, edata, gdata,
+        ntypes, etypes)
 end
 
 
@@ -138,36 +138,33 @@ end
 function Base.show(io::IO, ::MIME"text/plain", g::GNNHeteroGraph)
     if get(io, :compact, false)
         print(io, "GNNHeteroGraph($(g.num_nodes), $(g.num_edges))")
-    else # if the following block is indented the printing is ruined
-    print(io, "GNNHeteroGraph:
-  num_nodes: $((g.num_nodes...,))         
-  num_edges: $((g.num_edges...,))")
-  g.num_graphs > 1 && print(io, "\n    num_graphs = $(g.num_graphs)")
-  if !isempty(g.ndata)
-      print(io, "\n  ndata:")
-      for k in keys(g.ndata)
-          print(io, "\n    ", _str(k), "  =>  $(shortsummary(g.ndata[k]))")
-      end
-  end
-  if !isempty(g.edata)
-      print(io, "\n  edata:")
-      for k in keys(g.edata)
-          print(io, "\n    $k  =>  $(shortsummary(g.edata[k]))")
-      end
-  end
-  if !isempty(g.gdata)
-      print(io, "\n  gdata:")
-      print(io, "\n    ")
-      shortsummary(io, g.gdata)
-  end #else
-  end
+    else
+        print(io, "GNNHeteroGraph:\n num_nodes: $((g.num_nodes...,))\n num_edges: $((g.num_edges...,))")
+        g.num_graphs > 1 && print(io, "\n num_graphs: $(g.num_graphs)")
+        if !isempty(g.ndata)
+            print(io, "\n ndata:")
+            for k in keys(g.ndata)
+                print(io, "\n\t", _str(k), "  =>  $(shortsummary(g.ndata[k]))")
+            end
+        end
+        if !isempty(g.edata)
+            print(io, "\n edata:")
+            for k in keys(g.edata)
+                print(io, "\n\t$k  =>  $(shortsummary(g.edata[k]))")
+            end
+        end
+        if !isempty(g.gdata)
+            print(io, "\n gdata:\n\t")
+            shortsummary(io, g.gdata)
+        end
+    end
 end
 
 GNNHeteroGraph(data; kws...) = GNNHeteroGraph(Dict(data); kws...)
 
 _str(s::Symbol) = ":$s"
 _str(s) = "$s"
 
-MLUtils.numobs(g::GNNHeteroGraph) = g.num_graphs 
+MLUtils.numobs(g::GNNHeteroGraph) = g.num_graphs
 # MLUtils.getobs(g::GNNHeteroGraph, i) = getgraph(g, i)
 

test/GNNGraphs/gnnheterograph.jl

@@ -75,4 +75,24 @@ using Test
         @test hg.num_nodes == Dict(:A => 10, :B => 20)         
         @test hg.num_edges == Dict((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)
     end
+
+    @testset "show" begin
+        num_nodes = Dict(:A => 10, :B => 20);
+        edges1 = rand(1:num_nodes[:A], 20), rand(1:num_nodes[:B], 20)
+        edges2 = rand(1:num_nodes[:B], 30), rand(1:num_nodes[:A], 30)
+        eindex = ((:A, :rel1, :B) => edges1, (:B, :rel2, :A) => edges2)
+        ndata = Dict(:A => (x = rand(2, num_nodes[:A]), y = rand(3, num_nodes[:A])),:B => rand(10, num_nodes[:B]))
+        edata= Dict((:A, :rel1, :B) => (x = rand(2, 20), y = rand(3, 20)),(:B, :rel2, :A) => rand(10, 30))
+        hg1 = GraphNeuralNetworks.GNNHeteroGraph(eindex; num_nodes)
+        hg2 = GraphNeuralNetworks.GNNHeteroGraph(eindex; num_nodes, ndata,edata)
+        hg3 = GraphNeuralNetworks.GNNHeteroGraph(eindex; num_nodes, ndata)
+        @test sprint(show, hg1) == "GNNHeteroGraph(Dict(:A => 10, :B => 20), Dict((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30))"
+        @test sprint(show, hg2) == sprint(show, hg1)
+        @test sprint(show, MIME("text/plain"), hg1; context=:compact => true) == "GNNHeteroGraph(Dict(:A => 10, :B => 20), Dict((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30))"
+        @test sprint(show, MIME("text/plain"), hg2; context=:compact => true) == sprint(show, MIME("text/plain"), hg1;context=:compact => true)
+        @test sprint(show, MIME("text/plain"), hg1; context=:compact => false) == "GNNHeteroGraph:\n num_nodes: (:A => 10, :B => 20)\n num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)"
+        @test sprint(show, MIME("text/plain"), hg2; context=:compact => false) == "GNNHeteroGraph:\n num_nodes: (:A => 10, :B => 20)\n num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)\n ndata:\n\t:A  =>  (x = 2×10 Matrix{Float64}, y = 3×10 Matrix{Float64})\n\t:B  =>  x = 10×20 Matrix{Float64}\n edata:\n\t(:A, :rel1, :B)  =>  (x = 2×20 Matrix{Float64}, y = 3×20 Matrix{Float64})\n\t(:B, :rel2, :A)  =>  e = 10×30 Matrix{Float64}"
+        @test sprint(show, MIME("text/plain"), hg3; context=:compact => false) =="GNNHeteroGraph:\n num_nodes: (:A => 10, :B => 20)\n num_edges: ((:A, :rel1, :B) => 20, (:B, :rel2, :A) => 30)\n ndata:\n\t:A  =>  (x = 2×10 Matrix{Float64}, y = 3×10 Matrix{Float64})\n\t:B  =>  x = 10×20 Matrix{Float64}"
+        @test sprint(show, MIME("text/plain"), hg2; context=:compact => false) != sprint(show, MIME("text/plain"), hg3; context=:compact => false)
+    end
 end