[Misc] polish the python code (#6499)

python/dgl/graphbolt/dataset.py

@@ -32,17 +32,17 @@ class Task:
         raise NotImplementedError
 
     @property
-    def train_set(self) -> ItemSet or ItemSetDict:
+    def train_set(self) -> Union[ItemSet, ItemSetDict]:
         """Return the training set."""
         raise NotImplementedError
 
     @property
-    def validation_set(self) -> ItemSet or ItemSetDict:
+    def validation_set(self) -> Union[ItemSet, ItemSetDict]:
         """Return the validation set."""
         raise NotImplementedError
 
     @property
-    def test_set(self) -> ItemSet or ItemSetDict:
+    def test_set(self) -> Union[ItemSet, ItemSetDict]:
         """Return the test set."""
         raise NotImplementedError
 

python/dgl/graphbolt/item_sampler.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Mapping
 from functools import partial
-from typing import Callable, Iterator, Optional
+from typing import Callable, Iterator, Optional, Union
 
 import numpy as np
 import torch
@@ -200,7 +200,7 @@ class ItemSampler(IterDataPipe):
 
     Parameters
     ----------
-    item_set : ItemSet or ItemSetDict
+    item_set : Union[ItemSet, ItemSetDict]
         Data to be sampled.
     batch_size : int
         The size of each batch.
@@ -391,7 +391,7 @@ class ItemSampler(IterDataPipe):
 
     def __init__(
         self,
-        item_set: ItemSet or ItemSetDict,
+        item_set: Union[ItemSet, ItemSetDict],
         batch_size: int,
         minibatcher: Optional[Callable] = minibatcher_default,
         drop_last: Optional[bool] = False,
@@ -501,7 +501,7 @@ class DistributedItemSampler(ItemSampler):
 
     Parameters
     ----------
-    item_set : ItemSet or ItemSetDict
+    item_set : Union[ItemSet, ItemSetDict]
         Data to be sampled.
     batch_size : int
         The size of each batch.
@@ -627,7 +627,7 @@ class DistributedItemSampler(ItemSampler):
 
     def __init__(
         self,
-        item_set: ItemSet or ItemSetDict,
+        item_set: Union[ItemSet, ItemSetDict],
         batch_size: int,
         minibatcher: Optional[Callable] = minibatcher_default,
         drop_last: Optional[bool] = False,

python/dgl/graphbolt/subgraph_sampler.py

@@ -45,8 +45,8 @@ class SubgraphSampler(MiniBatchTransformer):
             seeds = minibatch.seed_nodes
         else:
             raise ValueError(
-                f"Invalid minibatch {minibatch}: Either 'node_pairs' or \
-                'seed_nodes' should have a value."
+                f"Invalid minibatch {minibatch}: Either `node_pairs` or "
+                "`seed_nodes` should have a value."
             )
         (
             minibatch.input_nodes,