Add getproperty function to TemporalSnapshotsGNNGraph (#297)

* Add `getproperty` for a vector of `Datastore`

* Add `getproperty` for a `TemporalSnapshotGNNGraph`

* Add `getproperty` docstring

* Remove comment

* Add test `getproperty`

* Add test `getproperty` tsg

* Fix typo

* Replace tgs with tg

* Add spaces

Co-authored-by: Carlo Lucibello <[email protected]>

* Remove docstring `Base.getproperty`

* Add more test `Datastore`

---------

Co-authored-by: Carlo Lucibello <[email protected]>

Author: aurorarossi

src/GNNGraphs/datastore.jl

@@ -105,6 +105,16 @@ function Base.getproperty(ds::DataStore, s::Symbol)
     end
 end
 
+function Base.getproperty(vds::Vector{DataStore}, s::Symbol)
+    if s === :_n
+        return [getn(ds) for ds in vds]
+    elseif s === :_data
+        return [getdata(ds) for ds in vds]
+    else
+        return [getdata(ds)[s] for ds in vds]
+    end
+end
+
 function Base.setproperty!(ds::DataStore, s::Symbol, x)
     @assert s != :_n "cannot set _n directly"
     @assert s != :_data "cannot set _data directly"

src/GNNGraphs/temporalsnapshotsgnngraph.jl

@@ -182,6 +182,20 @@ end
 #     return tg
 # end
 
+function Base.getproperty(tg::TemporalSnapshotsGNNGraph, prop::Symbol)
+    if prop ∈ fieldnames(TemporalSnapshotsGNNGraph)
+        return getfield(tg, prop)
+    elseif prop == :ndata
+        return [s.ndata for s in tg.snapshots]
+    elseif prop == :edata
+        return [s.edata for s in tg.snapshots]
+    elseif prop == :gdata
+        return [s.gdata for s in tg.snapshots]
+    else 
+        return [getproperty(s,prop) for s in tg.snapshots]
+    end
+end
+
 function Base.show(io::IO, tsg::TemporalSnapshotsGNNGraph)
     print(io, "TemporalSnapshotsGNNGraph($(tsg.num_snapshots)) with ")
     print_feature_t(io, tsg.tgdata)
@@ -205,7 +219,6 @@ function Base.show(io::IO, ::MIME"text/plain", tsg::TemporalSnapshotsGNNGraph)
     end
 end
 
-
 function print_feature_t(io::IO, feature)
     if !isempty(feature)
         if length(keys(feature)) == 1

test/GNNGraphs/datastore.jl

@@ -20,6 +20,11 @@ end
     @test_throws DimensionMismatch ds.z=rand(12)
     ds.z = [1:10;]
     @test ds.z == [1:10;]
+    vec = [DataStore(10, (:x => x,)), DataStore(10, (:x => x, :y => rand(2, 10)))]
+    @test vec.x == [x, x]
+    @test_throws KeyError vec.z 
+    @test vec._n == [10, 10]
+    @test vec._data == [Dict(:x => x), Dict(:x => x, :y => vec[2].y)]
 end
 
 @testset "map" begin
@@ -31,7 +36,7 @@ end
     @test_throws AssertionError ds2=map(x -> [x; x], ds)
 end
 
-@testset """getdata / getn""" begin
+@testset "getdata / getn" begin
     ds = DataStore(10, (:x => rand(10), :y => rand(2, 10)))
     @test getdata(ds) == getfield(ds, :_data)
     @test_throws KeyError ds.data

test/GNNGraphs/temporalsnapshotsgnngraph.jl

@@ -24,6 +24,16 @@ end
     @test tsg[[1,2]] == TemporalSnapshotsGNNGraph([10,10], [20,20], 2, snapshots[1:2], tsg.tgdata)
 end
 
+@testset "getproperty" begin
+    x = rand(10)
+    snapshots = [rand_graph(10, 20, ndata = x) for i in 1:5]
+    tsg = TemporalSnapshotsGNNGraph(snapshots)
+    @test tsg.tgdata == DataStore()
+    @test tsg.x == tsg.ndata.x == [x for i in 1:5]
+    @test_throws KeyError tsg.ndata.w
+    @test_throws ArgumentError tsg.w
+end
+
 @testset "add/remove_snapshot" begin
     snapshots = [rand_graph(10, 20) for i in 1:5]
     tsg = TemporalSnapshotsGNNGraph(snapshots)
@@ -91,5 +101,3 @@ end
     tsg.tgdata.x=rand(4)
     @test sprint(show,tsg) == "TemporalSnapshotsGNNGraph(5) with x: 4-element data"
 end
-
-#     @test sprint(show, MIME("text/plain"), rand_graph(10, 20); context=:compact => true) == "GNNGraph(10, 20) with no data"