adding tests + removing type annotation to avoid circular dep

test/data/test_dataset.py

@@ -1,3 +1,4 @@
+import pytest
 import torch
 from torch_sparse import SparseTensor
 
@@ -42,6 +43,20 @@ def test_in_memory_dataset():
     assert dataset[1].test_int == 2
     assert dataset[1].test_str == '2'
 
+    summary = dataset.summary()
+    num_nodes = torch.tensor([5., 10.])
+    assert summary.num_graphs == 2
+    assert summary.mean_num_nodes == pytest.approx(float(num_nodes.mean()))
+    assert summary.std_num_nodes == pytest.approx(float(num_nodes.std()))
+    assert summary.min_num_nodes == int(num_nodes.min())
+    assert summary.max_num_nodes == int(num_nodes.max())
+
+    num_edges = torch.tensor([4., 4.])
+    assert summary.mean_num_edges == pytest.approx(float(num_edges.mean()))
+    assert summary.std_num_edges == pytest.approx(float(num_edges.std()))
+    assert summary.min_num_edges == int(num_edges.min())
+    assert summary.max_num_edges == int(num_edges.max())
+
 
 def test_in_memory_sparse_tensor_dataset():
     x = torch.randn(11, 16)
@@ -118,6 +133,9 @@ def test_hetero_in_memory_dataset():
     assert (dataset[1]['paper', 'paper'].edge_index.tolist() == data2[
         'paper', 'paper'].edge_index.tolist())
 
+    summary = dataset.summary()
+    assert summary.num_graphs == 2
+
 
 def test_override_behaviour():
     class DS(Dataset):

torch_geometric/data/dataset.py

@@ -10,8 +10,9 @@
 import torch.utils.data
 from torch import Tensor
 
-from torch_geometric.data import Data, Summary
+from torch_geometric.data import Data
 from torch_geometric.data.makedirs import makedirs
+from torch_geometric.data.summary import Summary
 
 IndexType = Union[slice, Tensor, np.ndarray, Sequence]
 

torch_geometric/data/summary.py

@@ -2,16 +2,14 @@
 from tabulate import tabulate
 from tqdm import tqdm
 
-from torch_geometric.data import Dataset
-
 
 class Summary:
     r"""Summary of graph datasets
 
     Args:
         dataset (Dataset): :obj:`torch_geometric.data.Dataset`
     """
-    def __init__(self, dataset: Dataset):
+    def __init__(self, dataset):
         self.dataset_str = repr(dataset)
 
         def map_data(data):