Auto-document type hints (#9690)

https://pytorch-geometric--9690.org.readthedocs.build/en/9690/ --------- Co-authored-by: rusty1s <matthias.fey@tu-dortmund.de>
pyg-team · Oct 9, 2024 · 4019360 · 4019360
1 parent 3affeab
commit 4019360
Show file tree

Hide file tree

Showing 14 changed files with 105 additions and 112 deletions.
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -20,6 +20,7 @@
     'sphinx.ext.mathjax',
     'sphinx.ext.napoleon',
     'sphinx.ext.viewcode',
+    'sphinx_autodoc_typehints',
     'sphinx_copybutton',
     'nbsphinx',
     'pyg',
@@ -45,6 +46,9 @@
     'torch': ('https://pytorch.org/docs/main', None),
 }
 
+typehints_use_rtype = False
+typehints_defaults = 'comma'
+
 nbsphinx_thumbnails = {
     'tutorial/create_gnn':
     '_static/thumbnails/create_gnn.png',
@@ -78,5 +82,9 @@ def rst_jinja_render(app, _, source):
 
 
 def setup(app):
+    r"""Setup sphinx application."""
     app.connect('source-read', rst_jinja_render)
     app.add_js_file('js/version_alert.js')
+
+    # Do not drop type hints in signatures:
+    del app.events.listeners['autodoc-process-signature']
diff --git a/docs/source/modules/root.rst b/docs/source/modules/root.rst
@@ -24,6 +24,7 @@ Functions
 
 .. automodule:: torch_geometric.compile
     :members:
+    :exclude-members: compile
 
 .. automodule:: torch_geometric.debug
     :members:

diff --git a/torch_geometric/_compile.py b/torch_geometric/_compile.py
@@ -27,6 +27,11 @@ def compile(
     This function has the same signature as :meth:`torch.compile` (see
     `here <https://pytorch.org/docs/stable/generated/torch.compile.html>`__).
 
+    Args:
+        model: The model to compile.
+        *args: Additional arguments of :meth:`torch.compile`.
+        **kwargs: Additional keyword arguments of :meth:`torch.compile`.
+
     .. note::
         :meth:`torch_geometric.compile` is deprecated in favor of
         :meth:`torch.compile`.

diff --git a/torch_geometric/datasets/actor.py b/torch_geometric/datasets/actor.py
@@ -19,17 +19,15 @@ class Actor(InMemoryDataset):
     actor's Wikipedia.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in an
             :obj:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
-            transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before being saved to
+            disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

diff --git a/torch_geometric/datasets/airfrans.py b/torch_geometric/datasets/airfrans.py
@@ -46,26 +46,24 @@ class AirfRANS(InMemoryDataset):
         :obj:`torch_geometric.transforms.RadiusGraph` transform.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        task (str): The task to study (:obj:`"full"`, :obj:`"scarce"`,
+        root: Root directory where the dataset should be saved.
+        task: The task to study (:obj:`"full"`, :obj:`"scarce"`,
             :obj:`"reynolds"`, :obj:`"aoa"`) that defines the utilized training
             and test splits.
-        train (bool, optional): If :obj:`True`, loads the training dataset,
-            otherwise the test dataset. (default: :obj:`True`)
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        train: If :obj:`True`, loads the training dataset, otherwise the test
+            dataset.
+        transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        pre_filter (callable, optional): A function that takes in an
+            being saved to disk.
+        pre_filter: A function that takes in an
             :obj:`torch_geometric.data.Data` object and returns a boolean
             value, indicating whether the data object should be included in the
-            final dataset. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            final dataset.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

diff --git a/torch_geometric/datasets/airports.py b/torch_geometric/datasets/airports.py
@@ -14,22 +14,20 @@ class Airports(InMemoryDataset):
     and labels correspond to activity levels.
     Features are given by one-hot encoded node identifiers, as described in the
     `"GraLSP: Graph Neural Networks with Local Structural Patterns"
-    ` <https://arxiv.org/abs/1911.07675>`_ paper.
+    <https://arxiv.org/abs/1911.07675>`_ paper.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        name (str): The name of the dataset (:obj:`"USA"`, :obj:`"Brazil"`,
+        root: Root directory where the dataset should be saved.
+        name: The name of the dataset (:obj:`"USA"`, :obj:`"Brazil"`,
             :obj:`"Europe"`).
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
         pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
     """
     edge_url = ('https://github.com/leoribeiro/struc2vec/'
                 'raw/master/graph/{}-airports.edgelist')

diff --git a/torch_geometric/datasets/amazon.py b/torch_geometric/datasets/amazon.py
@@ -15,19 +15,16 @@ class Amazon(InMemoryDataset):
     map goods to their respective product category.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        name (str): The name of the dataset (:obj:`"Computers"`,
-            :obj:`"Photo"`).
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        root: Root directory where the dataset should be saved.
+        name: The name of the dataset (:obj:`"Computers"`, :obj:`"Photo"`).
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

diff --git a/torch_geometric/datasets/amazon_book.py b/torch_geometric/datasets/amazon_book.py
@@ -14,17 +14,16 @@ class AmazonBook(InMemoryDataset):
     No labels or features are provided.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.HeteroData` object and returns a
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in an
+            :class:`torch_geometric.data.HeteroData` object and returns a
             transformed version. The data object will be transformed before
-            every access. (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.HeteroData` object and returns a
+            every access.
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.HeteroData` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
     """
     url = ('https://raw.githubusercontent.com/gusye1234/LightGCN-PyTorch/'
            'master/data/amazon-book')

diff --git a/torch_geometric/datasets/amazon_products.py b/torch_geometric/datasets/amazon_products.py
@@ -14,17 +14,15 @@ class AmazonProducts(InMemoryDataset):
     containing products and its categories.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

diff --git a/torch_geometric/datasets/aminer.py b/torch_geometric/datasets/aminer.py
@@ -24,17 +24,16 @@ class AMiner(InMemoryDataset):
     truth labels for a subset of nodes.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.HeteroData` object and returns a
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.HeteroData` object and returns a
             transformed version. The data object will be transformed before
-            every access. (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.HeteroData` object and returns a
+            every access.
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.HeteroData` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
     """
 
     url = 'https://www.dropbox.com/s/1bnz8r7mofx0osf/net_aminer.zip?dl=1'

diff --git a/torch_geometric/datasets/aqsol.py b/torch_geometric/datasets/aqsol.py
@@ -30,25 +30,22 @@ class AQSOL(InMemoryDataset):
     the :class:`~torch_geometric.datasets.ZINC` dataset.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        split (str, optional): If :obj:`"train"`, loads the training dataset.
+        root: Root directory where the dataset should be saved.
+        split: If :obj:`"train"`, loads the training dataset.
             If :obj:`"val"`, loads the validation dataset.
             If :obj:`"test"`, loads the test dataset.
-            (default: :obj:`"train"`)
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
+            being saved to disk.
         pre_filter (callable, optional): A function that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a boolean
+            :class:`torch_geometric.data.Data` object and returns a boolean
             value, indicating whether the data object should be included in
-            the final dataset. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            the final dataset.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

diff --git a/torch_geometric/datasets/attributed_graph_dataset.py b/torch_geometric/datasets/attributed_graph_dataset.py
@@ -19,21 +19,19 @@ class AttributedGraphDataset(InMemoryDataset):
     <https://arxiv.org/abs/2009.00826>`_ paper.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        name (str): The name of the dataset (:obj:`"Wiki"`, :obj:`"Cora"`
+        root: Root directory where the dataset should be saved.
+        name: The name of the dataset (:obj:`"Wiki"`, :obj:`"Cora"`,
             :obj:`"CiteSeer"`, :obj:`"PubMed"`, :obj:`"BlogCatalog"`,
             :obj:`"PPI"`, :obj:`"Flickr"`, :obj:`"Facebook"`, :obj:`"Twitter"`,
             :obj:`"TWeibo"`, :obj:`"MAG"`).
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

diff --git a/torch_geometric/datasets/ba_multi_shapes.py b/torch_geometric/datasets/ba_multi_shapes.py
@@ -25,21 +25,19 @@ class BAMultiShapesDataset(InMemoryDataset):
     This dataset is pre-computed from the official implementation.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        pre_filter (callable, optional): A function that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a boolean
+            being saved to disk.
+        pre_filter: A function that takes in a
+            :class:`torch_geometric.data.Data` object and returns a boolean
             value, indicating whether the data object should be included in the
-            final dataset. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            final dataset.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

diff --git a/torch_geometric/datasets/ba_shapes.py b/torch_geometric/datasets/ba_shapes.py
@@ -30,15 +30,14 @@ class BAShapes(InMemoryDataset):
         :class:`torch_geometric.datasets.graph_generator.BAGraph` instead.
 
     Args:
-        connection_distribution (str, optional): Specifies how the houses
-            and the BA graph get connected. Valid inputs are :obj:`"random"`
+        connection_distribution: Specifies how the houses and the BA graph get
+            connected. Valid inputs are :obj:`"random"`
             (random BA graph nodes are selected for connection to the houses),
             and :obj:`"uniform"` (uniformly distributed BA graph nodes are
-            selected for connection to the houses). (default: :obj:`"random"`)
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+            selected for connection to the houses).
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
     """
     def __init__(
         self,