Repository: pyg-team/pytorch_geometric
Pull request #8750 (2024-01-12 08:18:46)

Re-factor `Inspector`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8750?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `41 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bc066b7`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/bc066b7b9b0ec16ef1110abdadd648e0d570aff6?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.56% compared to head [(`4ae2c10`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8750?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 23.90%.
> Report is 1 commits behind head on master.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8750?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/nn/conv/utils/inspector.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8750?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvdXRpbHMvaW5zcGVjdG9yLnB5) | 76.62% | [36 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8750?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |
| [torch\_geometric/nn/conv/message\_passing.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8750?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvbWVzc2FnZV9wYXNzaW5nLnB5) | 91.30% | [4 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8750?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |
| [torch\_geometric/nn/conv/propagate.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8750?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvcHJvcGFnYXRlLnB5) | 85.71% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8750?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           master    #8750       +/-   ##
===========================================
- Coverage   89.56%   23.90%   -65.66%     
===========================================
  Files         482      482               
  Lines       31113    31207       +94     
===========================================
- Hits        27865     7460    -20405     
- Misses       3248    23747    +20499     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8750?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/docs/source/modules/nn.rst b/docs/source/modules/nn.rst
index 29af7abc1..466c5b570 100644
--- a/docs/source/modules/nn.rst
+++ b/docs/source/modules/nn.rst
@@ -147,7 +147,7 @@ Finally, we added full support for customization of aggregations into the :class
 
 .. note::
 
-   You can read more about the :class:`torch_geometric.nn.aggr` package in this [blog post](https://medium.com/@pytorch_geometric/a-principled-approach-to-aggregations-983c086b10b3).
+   You can read more about the :class:`torch_geometric.nn.aggr` package in this `blog post <https://medium.com/@pytorch_geometric/a-principled-approach-to-aggregations-983c086b10b3>`__.
 
 .. autosummary::
    :nosignatures:
diff --git a/torch_geometric/nn/conv/message_passing.jinja b/torch_geometric/nn/conv/message_passing.jinja
index 6bdfd1bc2..5beb930af 100644
--- a/torch_geometric/nn/conv/message_passing.jinja
+++ b/torch_geometric/nn/conv/message_passing.jinja
@@ -47,7 +47,7 @@ class {{cls_name}}({{parent_cls_name}}):
 {%- if arg[-2:] not in ['_i', '_j'] %}
         {{arg}} = kwargs.{{arg}}
 {%- else %}
-        {{arg}}: {{collect_types[arg]}} = {% if collect_types[arg][:8] == 'Optional' %}None{% else %}init{% endif %}
+        {{arg}}: {{collect_types[arg]}} = {% if 'Optional' in collect_types[arg] %}None{% else %}init{% endif %}
         data = kwargs.{{arg[:-2]}}
         if isinstance(data, (tuple, list)):
             assert len(data) == 2
@@ -97,9 +97,9 @@ class {{cls_name}}({{parent_cls_name}}):
                 edge_type = value
             {% endif %}
 
-        {% if collect_types.get('edge_weight', 'Optional')[:8] != 'Optional' %}assert edge_weight is not None{% endif %}
-        {% if collect_types.get('edge_attr', 'Optional')[:8] != 'Optional' %}assert edge_attr is not None{% endif %}
-        {% if collect_types.get('edge_type', 'Optional')[:8] != 'Optional' %}assert edge_type is not None{% endif %}
+        {% if 'Optional' not in collect_types.get('edge_weight', 'Optional') %}assert edge_weight is not None{% endif %}
+        {% if 'Optional' not in collect_types.get('edge_attr', 'Optional') %}assert edge_attr is not None{% endif %}
+        {% if 'Optional' not in collect_types.get('edge_type', 'Optional') %}assert edge_type is not None{% endif %}
 
         index = edge_index_i
         size_i = size[1] if size[1] is not None else size[0]
@@ -122,7 +122,7 @@ class {{cls_name}}({{parent_cls_name}}):
 {%- if arg[-2:] not in ['_i', '_j'] %}
         {{arg}} = kwargs.{{arg}}
 {%- else %}
-        {{arg}}: {{edge_collect_types[arg]}} = {% if edge_collect_types[arg][:8] == 'Optional' %}None{% else %}init{% endif %}
+        {{arg}}: {{edge_collect_types[arg]}} = {% if 'Optional' in edge_collect_types %}None{% else %}init{% endif %}
         data = kwargs.{{arg[:-2]}}
         if isinstance(data, (tuple, list)):
             assert len(data) == 2
@@ -172,9 +172,9 @@ class {{cls_name}}({{parent_cls_name}}):
                 edge_type = value
             {% endif %}
 
-        {% if edge_collect_types.get('edge_weight', 'Optional')[:8] != 'Optional' %}assert edge_weight is not None{% endif %}
-        {% if edge_collect_types.get('edge_attr', 'Optional')[:8] != 'Optional' %}assert edge_attr is not None{% endif %}
-        {% if edge_collect_types.get('edge_type', 'Optional')[:8] != 'Optional' %}assert edge_type is not None{% endif %}
+        {% if 'Optional' not in edge_collect_types.get('edge_weight', 'Optional') %}assert edge_weight is not None{% endif %}
+        {% if 'Optional' not in edge_collect_types.get('edge_attr', 'Optional') %}assert edge_attr is not None{% endif %}
+        {% if 'Optional' not in edge_collect_types.get('edge_type', 'Optional') %}assert edge_type is not None{% endif %}
 
         index = edge_index_i
         size_i = size[1] if size[1] is not None else size[0]
diff --git a/torch_geometric/nn/conv/message_passing.py b/torch_geometric/nn/conv/message_passing.py
index 578ef991b..1db2b4651 100644
--- a/torch_geometric/nn/conv/message_passing.py
+++ b/torch_geometric/nn/conv/message_passing.py
@@ -1,8 +1,8 @@
 import inspect
 import os.path as osp
 import random
-import re
 import warnings
+from abc import abstractmethod
 from inspect import Parameter
 from typing import (
     Any,
@@ -13,8 +13,8 @@ from typing import (
     OrderedDict,
     Set,
     Tuple,
+    Type,
     Union,
-    get_type_hints,
 )
 
 import torch
@@ -23,13 +23,11 @@ from torch.utils.hooks import RemovableHandle
 
 from torch_geometric.nn.aggr import Aggregation
 from torch_geometric.nn.conv.propagate import (
-    find_parenthesis_content,
     module_from_template,
     type_hint_to_str,
 )
 from torch_geometric.nn.conv.utils.inspector import Inspector
 from torch_geometric.nn.conv.utils.jit import class_from_module_repr
-from torch_geometric.nn.conv.utils.typing import sanitize, split_types_repr
 from torch_geometric.nn.resolver import aggregation_resolver as aggr_resolver
 from torch_geometric.typing import Adj, Size, SparseTensor
 from torch_geometric.utils import (
@@ -140,20 +138,19 @@ class MessagePassing(torch.nn.Module):
         self.decomposed_layers = decomposed_layers
 
         # Collect attribute names requested in message passing hooks:
-        self.inspector = Inspector(self)
-        self.inspector.inspect(self.message)
-        self.inspector.inspect(self.aggregate, pop_first=True)
-        self.inspector.params['aggregate'].pop('aggr', None)
-        self.inspector.inspect(self.message_and_aggregate, pop_first=True)
-        self.inspector.inspect(self.update, pop_first=True)
-        self.inspector.inspect(self.edge_update)
-
-        self._user_args: Set[str] = self.inspector.keys(
-            ['message', 'aggregate', 'update']).difference(self.special_args)
-        self._fused_user_args: Set[str] = self.inspector.keys(
-            ['message_and_aggregate', 'update']).difference(self.special_args)
-        self._edge_user_args: Set[str] = self.inspector.keys(
-            ['edge_update']).difference(self.special_args)
+        self.inspector = Inspector(self.__class__)
+        self.inspector.inspect_signature(self.message)
+        self.inspector.inspect_signature(self.aggregate, exclude=[0, 'aggr'])
+        self.inspector.inspect_signature(self.message_and_aggregate, [0])
+        self.inspector.inspect_signature(self.update, exclude=[0])
+        self.inspector.inspect_signature(self.edge_update)
+
+        self._user_args: List[str] = self.inspector.get_flat_param_names(
+            ['message', 'aggregate', 'update'], exclude=self.special_args)
+        self._fused_user_args: List[str] = self.inspector.get_flat_param_names(
+            ['message_and_aggregate', 'update'], exclude=self.special_args)
+        self._edge_user_args: List[str] = self.inspector.get_param_names(
+            'edge_update', exclude=self.special_args)
 
         # Support for "fused" message passing:
         self.fuse = self.inspector.implements('message_and_aggregate')
@@ -182,56 +179,34 @@ class MessagePassing(torch.nn.Module):
         if getattr(self, 'jit_on_init', False):
             root_dir = osp.dirname(osp.realpath(__file__))
             prop_types, prop_return_type = self._get_propagate_types()
+            prop_types = {
+                param.name: type_hint_to_str(param.type)
+                for param in prop_types.values()
+            }
+            prop_return_type = type_hint_to_str(prop_return_type)
+            collect_types = self.inspector.get_flat_params(
+                ['message', 'aggregate', 'update'])
+            collect_types = {
+                param.name: type_hint_to_str(param.type)
+                for param in collect_types
+            }
             module = module_from_template(
                 module_name=f'{self.__module__}_propagate',
                 module=self.__module__,
                 template_path=osp.join(root_dir, 'propagate.jinja'),
                 propagate_types=prop_types,
                 propagate_return_type=prop_return_type,
-                collect_types=self.inspector.types(
-                    ['message', 'aggregate', 'update']),
-                message_args=self.inspector.keys(['message']),
-                aggregate_args=self.inspector.keys(['aggregate']),
-                message_and_aggregate_args=self.inspector.keys(
-                    ['message_and_aggregate']),
-                update_args=self.inspector.keys(['update']),
+                collect_types=collect_types,
+                message_args=self.inspector.get_param_names('message'),
+                aggregate_args=self.inspector.get_param_names('aggregate'),
+                message_and_aggregate_args=self.inspector.get_param_names(
+                    'message_and_aggregate'),
+                update_args=self.inspector.get_param_names('update'),
             )
 
             self.propagate_jit = module.propagate
             self._collect_jit = module._collect
 
-    def _get_propagate_types(self) -> Tuple[Dict[str, str], str]:
-        # Parse `propagate_types` and generate an efficient `propagate` method:
-        if hasattr(self, 'propagate_type'):
-            assert isinstance(self.propagate_type, dict)
-            propagate_types = {
-                name: type_hint_to_str(type_hint)
-                for name, type_hint in self.propagate_type.items()
-            }
-        else:
-            source = inspect.getsource(self.forward)
-            match = find_parenthesis_content(source, prefix='propagate_type:')
-            if match is not None:
-                propagate_types = dict(
-                    [re.split(r'\s*:\s*', t) for t in split_types_repr(match)])
-            else:
-                match = find_parenthesis_content(source, 'self.propagate')
-                if match is None:  # No `self.propagate` call:
-                    return
-                args = [  # Find all keyword argument names:
-                    arg[:arg.find('=')].strip() for arg in match.split(',')
-                    if arg.find('=') >= 0
-                ]
-                propagate_types = {
-                    arg: 'Tensor'
-                    for arg in args if arg not in {'edge_index', 'size'}
-                }
-
-        propagate_return_type = type_hint_to_str(
-            get_type_hints(self.update).get('return', Tensor))
-
-        return propagate_types, propagate_return_type
-
     def reset_parameters(self) -> None:
         r"""Resets all learnable parameters of the module."""
         if self.aggr_module is not None:
@@ -508,7 +483,7 @@ class MessagePassing(torch.nn.Module):
             coll_dict = self._collect(self._fused_user_args, edge_index,
                                       mutable_size, kwargs)
 
-            msg_aggr_kwargs = self.inspector.distribute(
+            msg_aggr_kwargs = self.inspector.collect_param_data(
                 'message_and_aggregate', coll_dict)
             for hook in self._message_and_aggregate_forward_pre_hooks.values():
                 res = hook(self, (edge_index, msg_aggr_kwargs))
@@ -520,7 +495,8 @@ class MessagePassing(torch.nn.Module):
                 if res is not None:
                     out = res
 
-            update_kwargs = self.inspector.distribute('update', coll_dict)
+            update_kwargs = self.inspector.collect_param_data(
+                'update', coll_dict)
             out = self.update(out, **update_kwargs)
 
         else:  # Otherwise, run both functions in separation.
@@ -541,7 +517,8 @@ class MessagePassing(torch.nn.Module):
                 coll_dict = self._collect(self._user_args, edge_index,
                                           mutable_size, kwargs)
 
-                msg_kwargs = self.inspector.distribute('message', coll_dict)
+                msg_kwargs = self.inspector.collect_param_data(
+                    'message', coll_dict)
                 for hook in self._message_forward_pre_hooks.values():
                     res = hook(self, (msg_kwargs, ))
                     if res is not None:
@@ -553,11 +530,12 @@ class MessagePassing(torch.nn.Module):
                         out = res
 
                 if self.explain:
-                    explain_msg_kwargs = self.inspector.distribute(
+                    explain_msg_kwargs = self.inspector.collect_param_data(
                         'explain_message', coll_dict)
                     out = self.explain_message(out, **explain_msg_kwargs)
 
-                aggr_kwargs = self.inspector.distribute('aggregate', coll_dict)
+                aggr_kwargs = self.inspector.collect_param_data(
+                    'aggregate', coll_dict)
                 for hook in self._aggregate_forward_pre_hooks.values():
                     res = hook(self, (aggr_kwargs, ))
                     if res is not None:
@@ -570,7 +548,8 @@ class MessagePassing(torch.nn.Module):
                     if res is not None:
                         out = res
 
-                update_kwargs = self.inspector.distribute('update', coll_dict)
+                update_kwargs = self.inspector.collect_param_data(
+                    'update', coll_dict)
                 out = self.update(out, **update_kwargs)
 
                 if decomposed_layers > 1:
@@ -618,6 +597,7 @@ class MessagePassing(torch.nn.Module):
         return self.aggr_module(inputs, index, ptr=ptr, dim_size=dim_size,
                                 dim=self.node_dim)
 
+    @abstractmethod
     def message_and_aggregate(
         self,
         adj_t: Adj,
@@ -666,7 +646,8 @@ class MessagePassing(torch.nn.Module):
         coll_dict = self._collect(self._edge_user_args, edge_index,
                                   mutable_size, kwargs)
 
-        edge_kwargs = self.inspector.distribute('edge_update', coll_dict)
+        edge_kwargs = self.inspector.collect_param_data(
+            'edge_update', coll_dict)
         out = self.edge_update(**edge_kwargs)
 
         for hook in self._edge_update_forward_hooks.values():
@@ -676,6 +657,7 @@ class MessagePassing(torch.nn.Module):
 
         return out
 
+    @abstractmethod
     def edge_update(self) -> Tensor:
         r"""Computes or updates features for each edge in the graph.
         This function can take any argument as input which was initially passed
@@ -704,9 +686,9 @@ class MessagePassing(torch.nn.Module):
             methods = ['message', 'aggregate', 'update']
 
         self._explain = explain
-        self.inspector.inspect(self.explain_message, pop_first=True)
-        self._user_args = self.inspector.keys(methods).difference(
-            self.special_args)
+        self.inspector.inspect_signature(self.explain_message, exclude=[0])
+        self._user_args = self.inspector.get_flat_param_names(
+            methods, exclude=self.special_args)
 
     def explain_message(self, inputs: Tensor, size_i: int) -> Tensor:
         # NOTE Replace this method in custom explainers per message-passing
@@ -886,6 +868,18 @@ class MessagePassing(torch.nn.Module):
 
     # TorchScript Support #####################################################
 
+    def _get_propagate_types(self) -> Tuple[Dict[str, Parameter], Type]:
+        param_dict = self.inspector.get_params_from_method_call(
+            'propagate', exclude=[0, 'edge_index', 'size'])
+        return_type = self.inspector.get_signature('update').return_type
+        return param_dict, return_type
+
+    def _get_edge_updater_types(self) -> Tuple[Dict[str, Parameter], Type]:
+        param_dict = self.inspector.get_params_from_method_call(
+            'edge_updater', exclude=[0, 'edge_index'])
+        return_type = self.inspector.get_signature('edge_update').return_type
+        return param_dict, return_type
+
     @torch.jit.unused
     def jittable(self, typing: Optional[str] = None) -> 'MessagePassing':
         r"""Analyzes the :class:`MessagePassing` instance and produces a new
@@ -901,46 +895,38 @@ class MessagePassing(torch.nn.Module):
 
         from jinja2 import Template
 
-        source = inspect.getsource(self.__class__)
+        source = self.inspector.source
 
         # Find and parse `propagate()` types to format `{arg1: type1, ...}`.
         prop_types, prop_return_type = self._get_propagate_types()
+        prop_types = {
+            param.name: type_hint_to_str(param.type)
+            for param in prop_types.values()
+        }
+        prop_return_type = type_hint_to_str(prop_return_type)
 
         # Find and parse `edge_updater` types to format `{arg1: type1, ...}`.
-        if 'edge_update' in self.__class__.__dict__.keys():
-            if hasattr(self, 'edge_updater_type'):
-                assert isinstance(self.edge_updater_type, dict)
-                edge_updater_types = {
-                    k: sanitize(str(v))
-                    for k, v in self.edge_updater_type.items()
-                }
-            else:
-                match = find_parenthesis_content(source, 'edge_updater_type:')
-                if match is None:
-                    raise TypeError(
-                        'TorchScript support requires the definition of the '
-                        'types passed to `edge_updater()`. Please specify '
-                        'them via\n\n edge_updater_type = {"arg1": type1, '
-                        '"arg2": type2, ... }\n\n or via\n\n'
-                        '# edge_updater_type: (arg1: type1, arg2: type2, ...)'
-                        '\n\ninside the `MessagePassing` module.')
-                edge_updater_types = dict(
-                    [re.split(r'\s*:\s*', t) for t in split_types_repr(match)])
-        else:
-            edge_updater_types = {}
-
-        type_hints = get_type_hints(self.__class__.edge_update)
-        edge_updater_return_type = type_hints.get('return', 'Tensor')
-        if str(edge_updater_return_type)[:6] == '<class':
-            edge_updater_return_type = edge_updater_return_type.__name__
+        edge_types, edge_return_type = self._get_edge_updater_types()
+        edge_types = {
+            param.name: type_hint_to_str(param.type)
+            for param in edge_types.values()
+        }
+        edge_return_type = type_hint_to_str(edge_return_type)
 
         # Parse `_collect()` types to format `{arg:1, type1, ...}`.
-        collect_types = self.inspector.types(
+        collect_types = self.inspector.get_flat_params(
             ['message', 'aggregate', 'update'])
-
-        # Parse `_collect()` types to format `{arg:1, type1, ...}`,
-        # specific to the argument used for edge updates.
-        edge_collect_types = self.inspector.types(['edge_update'])
+        collect_types = {
+            param.name: type_hint_to_str(param.type)
+            for param in collect_types
+        }
+
+        # Parse `_edge_collect()` types to format `{arg:1, type1, ...}`.
+        edge_collect_types = self.inspector.get_params('edge_update')
+        edge_collect_types = {
+            param.name: type_hint_to_str(param.type)
+            for param in edge_collect_types
+        }
 
         # Collect `forward()` function in case it is overloaded. This is
         # necessary since TorchScript cannot handle inheritance of overloaded
@@ -970,14 +956,15 @@ class MessagePassing(torch.nn.Module):
             collect_types=collect_types,
             user_args=self._user_args,
             edge_user_args=self._edge_user_args,
-            msg_args=self.inspector.keys(['message']),
-            aggr_args=self.inspector.keys(['aggregate']),
-            msg_and_aggr_args=self.inspector.keys(['message_and_aggregate']),
-            update_args=self.inspector.keys(['update']),
+            msg_args=self.inspector.get_param_names('message'),
+            aggr_args=self.inspector.get_param_names('aggregate'),
+            msg_and_aggr_args=self.inspector.get_param_names(
+                'message_and_aggregate'),
+            update_args=self.inspector.get_param_names('update'),
             edge_collect_types=edge_collect_types,
-            edge_update_args=self.inspector.keys(['edge_update']),
-            edge_updater_types=edge_updater_types,
-            edge_updater_return_type=edge_updater_return_type,
+            edge_update_args=self.inspector.get_param_names('edge_update'),
+            edge_updater_types=edge_types,
+            edge_updater_return_type=edge_return_type,
             forward_repr=forward_repr,
         )
         # Instantiate a class from the rendered JIT module representation.
diff --git a/torch_geometric/nn/conv/propagate.py b/torch_geometric/nn/conv/propagate.py
index 2dc629e71..553533b4c 100644
--- a/torch_geometric/nn/conv/propagate.py
+++ b/torch_geometric/nn/conv/propagate.py
@@ -3,7 +3,7 @@ import os
 import os.path as osp
 import re
 import sys
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from jinja2 import Template
 
@@ -37,6 +37,14 @@ def module_from_template(
 
 
 def type_hint_to_str(type_hint: Any) -> str:
+    if (getattr(type_hint, '__origin__', None) == Union
+            and len(getattr(type_hint, '__args__', [])) == 2):
+        arg1, arg2 = type_hint.__args__
+        if arg1 is type(None):
+            return f'typing.Optional[{type_hint_to_str(arg2)}]'
+        if arg2 is type(None):
+            return f'typing.Optional[{type_hint_to_str(arg1)}]'
+
     type_repr = str(type_hint)
     type_repr = re.sub(r'<class \'(.*)\'>', r'\1', type_repr)
     return type_repr
diff --git a/torch_geometric/nn/conv/utils/inspector.py b/torch_geometric/nn/conv/utils/inspector.py
index e5d42b77a..176510434 100644
--- a/torch_geometric/nn/conv/utils/inspector.py
+++ b/torch_geometric/nn/conv/utils/inspector.py
@@ -1,86 +1,353 @@
 import inspect
 import re
-from collections import OrderedDict
-from typing import Any, Callable, Dict, List, Optional, Set
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, Type, Union
 
-from torch_geometric.nn.conv.utils.typing import parse_types
+from torch import Tensor
+
+
+class Parameter(NamedTuple):
+    name: str
+    type: Union[Type, str]
+    default: Any
+
+
+class Signature(NamedTuple):
+    param_dict: Dict[str, Parameter]
+    return_type: Type
 
 
 class Inspector:
-    def __init__(self, base_class: Any):
-        self.base_class: Any = base_class
-        self.params: Dict[str, Dict[str, Any]] = {}
-
-    def inspect(self, func: Callable,
-                pop_first: bool = False) -> Dict[str, Any]:
-        params = inspect.signature(func).parameters
-        params = OrderedDict(params)
-        if pop_first:
-            params.popitem(last=False)
-        self.params[func.__name__] = params
-
-    def keys(self, func_names: Optional[List[str]] = None) -> Set[str]:
-        keys = []
-        for func in func_names or list(self.params.keys()):
-            keys += self.params[func].keys()
-        return set(keys)
-
-    def _implements(self, cls, func_name: str) -> bool:
-        if cls.__name__ == 'MessagePassing':
-            return False
-        if func_name in cls.__dict__.keys():
-            return True
-        return any(self._implements(c, func_name) for c in cls.__bases__)
+    r"""Inspects a given class and collects information about its instance
+    methods.
+
+    Args:
+        cls (Type): The class to inspect.
+    """
+    def __init__(self, cls: Type):
+        self._cls = cls
+        self._signature_dict: Dict[str, Signature] = {}
+        self._source: Optional[str] = None
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self._cls.__name__})'
 
     def implements(self, func_name: str) -> bool:
-        return self._implements(self.base_class.__class__, func_name)
-
-    def types(self, func_names: Optional[List[str]] = None) -> Dict[str, str]:
-        out: Dict[str, str] = {}
-        for func_name in func_names or list(self.params.keys()):
-            func = getattr(self.base_class, func_name)
-            arg_types = parse_types(func)[0][0]
-            for key in self.params[func_name].keys():
-                if key in out and out[key] != arg_types[key]:
-                    raise ValueError(
-                        (f'Found inconsistent types for argument {key}. '
-                         f'Expected type {out[key]} but found type '
-                         f'{arg_types[key]}.'))
-                out[key] = arg_types[key]
-        return out
-
-    def distribute(self, func_name, kwargs: Dict[str, Any]):
-        out = {}
-        for key, param in self.params[func_name].items():
-            data = kwargs.get(key, inspect.Parameter.empty)
-            if data is inspect.Parameter.empty:
-                if param.default is inspect.Parameter.empty:
-                    raise TypeError(f'Required parameter {key} is empty.')
-                data = param.default
-            out[key] = data
-        return out
-
-
-def func_header_repr(func: Callable, keep_annotation: bool = True) -> str:
-    source = inspect.getsource(func)
-    signature = inspect.signature(func)
-
-    if keep_annotation:
-        return ''.join(re.split(r'(\).*?:.*?\n)', source,
-                                maxsplit=1)[:2]).strip()
-
-    params_repr = ['self']
-    for param in signature.parameters.values():
-        params_repr.append(param.name)
-        if param.default is not inspect.Parameter.empty:
-            params_repr[-1] += f'={param.default}'
-
-    return f'def {func.__name__}({", ".join(params_repr)}):'
-
-
-def func_body_repr(func: Callable, keep_annotation: bool = True) -> str:
-    source = inspect.getsource(func)
-    body_repr = re.split(r'\).*?:.*?\n', source, maxsplit=1)[1]
-    if not keep_annotation:
-        body_repr = re.sub(r'\s*# type:.*\n', '', body_repr)
-    return body_repr
+        r"""Returns :obj:`True` in case the inspected class implements the
+        :obj:`func_name` method.
+
+        Args:
+            func_name (str): The function name to check for existence.
+        """
+        func = getattr(self._cls, func_name, None)
+        if not callable(func):
+            return False
+        return not getattr(func, '__isabstractmethod__', False)
+
+    # Inspecting Method Signatures ############################################
+
+    def inspect_signature(
+        self,
+        func: Union[Callable, str],
+        exclude: Optional[List[Union[str, int]]] = None,
+    ) -> Signature:
+        r"""Inspects the function signature of :obj:`func` and returns a tuple
+        of parameter types and return type.
+
+        Args:
+            func (callabel or str): The function.
+            exclude (list[int or str]): A list of parameters to exclude, either
+                given by their name or index. (default: :obj:`None`)
+        """
+        if isinstance(func, str):
+            func = getattr(self._cls, func)
+
+        if func.__name__ in self._signature_dict:
+            return self._signature_dict[func.__name__]
+
+        signature = inspect.signature(func)
+        params = [p for p in signature.parameters.values() if p.name != 'self']
+
+        param_dict: Dict[str, Parameter] = {}
+        for i, param in enumerate(params):
+            if exclude is not None and (i in exclude or param.name in exclude):
+                continue
+
+            param_type = param.annotation
+            param_dict[param.name] = Parameter(
+                name=param.name,
+                # Mimic TorchScript which auto-infers `Tensor` on empty types:
+                type=Tensor if param_type is inspect._empty else param_type,
+                default=param.default,
+            )
+
+        return_type = signature.return_annotation
+        return_type = Tensor if return_type is inspect._empty else return_type
+
+        self._signature_dict[func.__name__] = Signature(
+            param_dict, return_type)
+
+        return self._signature_dict[func.__name__]
+
+    def get_signature(
+        self,
+        func: Union[Callable, str],
+        exclude: Optional[List[str]] = None,
+    ) -> Signature:
+        r"""Returns the function signature of the inspected function
+        :obj:`func`.
+
+        Args:
+            func (callabel or str): The function.
+            exclude (list[str], optional): The parameter names to exclude.
+                (default: :obj:`None`)
+        """
+        func_name = func if isinstance(func, str) else func.__name__
+        signature = self._signature_dict.get(func_name)
+        if signature is None:
+            raise IndexError(f"Could not access signature for function "
+                             f"'{func_name}'. Did you forget to inspect it?")
+
+        if exclude is None:
+            return signature
+
+        param_dict = signature.param_dict
+        param_dict = {k: v for k, v in param_dict.items() if k not in exclude}
+        return Signature(param_dict, signature.return_type)
+
+    def get_params(
+        self,
+        func: Union[Callable, str],
+        exclude: Optional[List[str]] = None,
+    ) -> List[Parameter]:
+        r"""Returns the parameters of the inspected function :obj:`func`.
+
+        Args:
+            func (str or callable): The function.
+            exclude (list[str], optional): The parameter names to exclude.
+                (default: :obj:`None`)
+        """
+        signature = self.get_signature(func, exclude)
+        return [param for param in signature.param_dict.values()]
+
+    def get_flat_params(
+        self,
+        funcs: List[Union[Callable, str]],
+        exclude: Optional[List[str]] = None,
+    ) -> List[Parameter]:
+        r"""Returns the union of parameters of all inspected functions in
+        :obj:`funcs`.
+
+        Args:
+            funcs (list[str or callable]): The functions.
+            exclude (list[str], optional): The parameter names to exclude.
+                (default: :obj:`None`)
+        """
+        out_dict: Dict[str, Parameter] = {}
+        for func in funcs:
+            params = self.get_params(func, exclude)
+            for param in params:
+                expected = out_dict.get(param.name)
+                if expected is not None and param.type != expected.type:
+                    raise ValueError(f"Found inconsistent types for argument "
+                                     f"'{param.name}'. Expected type "
+                                     f"'{expected.type}' but found type "
+                                     f"'{param.type}'.")
+
+                if expected is not None and param.default != expected.default:
+                    if (param.default is not inspect._empty
+                            and expected.default is not inspect._empty):
+                        raise ValueError(f"Found inconsistent defaults for "
+                                         f"argument '{param.name}'. Expected "
+                                         f"'{expected.default}'  but found "
+                                         f"'{param.default}'.")
+
+                    default = expected.default
+                    if default is inspect._empty:
+                        default = param.default
+
+                    out_dict[param.name] = Parameter(  #
+                        param.name, param.type, default)
+
+                if expected is None:
+                    out_dict[param.name] = param
+
+        return list(out_dict.values())
+
+    def get_param_names(
+        self,
+        func: Union[Callable, str],
+        exclude: Optional[List[str]] = None,
+    ) -> List[str]:
+        r"""Returns the parameter names of the inspected function :obj:`func`.
+
+        Args:
+            func (str or callable): The function.
+            exclude (list[str], optional): The parameter names to exclude.
+                (default: :obj:`None`)
+        """
+        return [param.name for param in self.get_params(func, exclude)]
+
+    def get_flat_param_names(
+        self,
+        funcs: List[Union[Callable, str]],
+        exclude: Optional[List[str]] = None,
+    ) -> List[str]:
+        r"""Returns the union of parameter names of all inspected functions in
+        :obj:`funcs`.
+
+        Args:
+            funcs (list[str or callable]): The functions.
+            exclude (list[str], optional): The parameter names to exclude.
+                (default: :obj:`None`)
+        """
+        return [param.name for param in self.get_flat_params(funcs, exclude)]
+
+    def collect_param_data(
+        self,
+        func: Union[Callable, str],
+        kwargs: Dict[str, Any],
+    ) -> Dict[str, Any]:
+        r"""Collects the input data of the inspected function :obj:`func`
+        according to its function signature from a data blob.
+
+        Args:
+            func (callabel or str): The function.
+            kwargs (dict[str, Any]): The data blob which may serve as inputs.
+        """
+        out_dict: Dict[str, Any] = {}
+        for param in self.get_params(func):
+            if param.name not in kwargs:
+                if param.default is inspect._empty:
+                    raise TypeError(f"Parameter '{param.name}' is required")
+                out_dict[param.name] = param.default
+            else:
+                out_dict[param.name] = kwargs[param.name]
+        return out_dict
+
+    # Inspecting Method Bodies ################################################
+
+    @property
+    def source(self) -> str:
+        r"""Returns the source code of :obj:`cls`."""
+        if self._source is not None:
+            return self._source
+        self._source = inspect.getsource(self._cls)
+        return self._source
+
+    def get_params_from_method_call(
+        self,
+        func: Union[Callable, str],
+        exclude: Optional[List[Union[int, str]]] = None,
+    ) -> Dict[str, Parameter]:
+        r"""Parses a method call of :obj:`func` and returns its keyword
+        arguments.
+
+        .. note::
+            The method is required to be called via keyword arguments in case
+            type annotations are not found.
+
+        Args:
+            func (callabel or str): The function.
+            exclude (list[int or str]): A list of parameters to exclude, either
+                given by their name or index. (default: :obj:`None`)
+        """
+        func_name = func if isinstance(func, str) else func.__name__
+
+        # Three ways to specify the parameters of an unknown function header:
+        # 1. Defined as class attributes in `{func_name}_type`.
+        # 2. Defined via type annotations in `# {func_name}_type: (...)`.
+        # 3. Defined via parsing of the function call.
+
+        # (1) Find class attribute:
+        if hasattr(self._cls, f'{func_name}_type'):
+            type_dict = getattr(self._cls, f'{func_name}_type')
+            if not isinstance(type_dict, dict):
+                raise ValueError(f"'{func_name}_type' is expected to be a "
+                                 f"dictionary (got '{type(type_dict)}')")
+            return {
+                name: Parameter(name, param_type, inspect._empty)
+                for name, param_type in type_dict.items()
+            }
+
+        # (2) Find type annotation:
+        match = find_parenthesis_content(self.source, f'{func_name}_type:')
+        if match is not None:
+            param_dict: Dict[str, Parameter] = {}
+            for type_repr in split(match, sep=','):
+                name_and_type = re.split(r'\s*:\s*', type_repr)
+                if len(name_and_type) != 2:
+                    raise ValueError(f"Could not parse the argument "
+                                     f"'{type_repr}' of the "
+                                     f"'{func_name}_type' annotation")
+                name, param_type = name_and_type[0], name_and_type[1]
+                param_dict[name] = Parameter(name, param_type, inspect._empty)
+            return param_dict
+
+        # (3) Parse the function call:
+        match = find_parenthesis_content(self.source, f'self.{func_name}')
+        if match is not None:
+            param_dict: Dict[str, Parameter] = {}
+            for i, type_repr in enumerate(split(match, sep=',')):
+                if exclude is not None and i in exclude:
+                    continue
+                name_and_content = re.split(r'\s*=\s*', type_repr)
+                if len(name_and_content) != 2:
+                    raise ValueError(f"Could not parse the keyword argument "
+                                     f"'{type_repr}' of the "
+                                     f"'self.{func_name}' call")
+                name = name_and_content[0]
+                if exclude is not None and name in exclude:
+                    continue
+                param_dict[name] = Parameter(name, Tensor, inspect._empty)
+            return param_dict
+
+        return {}  # (4) No function call found:
+
+
+def find_parenthesis_content(source: str, prefix: str) -> Optional[str]:
+    r"""Returns the content of :obj:`{prefix}.*(...)` within :obj:`source`."""
+    match = re.search(prefix, source)
+    if match is None:
+        return
+
+    offset = source[match.start():].find('(')
+    if offset < 0:
+        return
+
+    source = source[match.start() + offset:]
+
+    depth = 0
+    for end, char in enumerate(source):
+        if char == '(':
+            depth += 1
+        if char == ')':
+            depth -= 1
+        if depth == 0:
+            content = source[1:end]
+            # Properly handle line breaks and multiple white-spaces:
+            content = content.replace('\n', ' ')
+            content = content.replace('#', ' ')
+            content = re.sub(' +', ' ', content)
+            content = content.strip()
+            return content
+
+
+def split(content: str, sep: str) -> List[str]:
+    r"""Splits :obj:`content` based on :obj:`sep`.
+    :obj:`sep` inside parentheses or square brackets are ignored.
+    """
+    assert len(sep) == 1
+    outs: List[str] = []
+
+    start = depth = 0
+    for end, char in enumerate(content):
+        if char == '[' or char == '(':
+            depth += 1
+        elif char == ']' or char == ')':
+            depth -= 1
+        elif char == sep and depth == 0:
+            outs.append(content[start:end].strip())
+            start = end + 1
+    if start != len(content):  # Respect dangling `sep`:
+        outs.append(content[start:].strip())
+    return outs


Test patch:
diff --git a/test/nn/conv/test_propagate.py b/test/nn/conv/test_propagate.py
index ec662c7f3..3f3f12f96 100644
--- a/test/nn/conv/test_propagate.py
+++ b/test/nn/conv/test_propagate.py
@@ -10,8 +10,11 @@ def test_propagate():
 
     old_propagate = conv.__class__.propagate
     old_collect = conv.__class__._collect
-    conv.__class__.propagate = conv.propagate_jit
-    conv.__class__._collect = conv._collect_jit
-    torch.jit.script(conv)
-    conv.__class__.propagate = old_propagate
-    conv.__class__._collect = old_collect
+
+    try:
+        conv.__class__.propagate = conv.propagate_jit
+        conv.__class__._collect = conv._collect_jit
+        torch.jit.script(conv)
+    finally:
+        conv.__class__.propagate = old_propagate
+        conv.__class__._collect = old_collect
diff --git a/test/nn/conv/utils/test_inspector.py b/test/nn/conv/utils/test_inspector.py
new file mode 100644
index 000000000..55bd66523
--- /dev/null
+++ b/test/nn/conv/utils/test_inspector.py
@@ -0,0 +1,106 @@
+import inspect
+
+import torch
+from torch import Tensor
+
+from torch_geometric.nn import GATConv, SAGEConv
+from torch_geometric.nn.conv.utils.inspector import (
+    Inspector,
+    Parameter,
+    Signature,
+)
+
+
+def test_inspector_sage_conv() -> None:
+    inspector = Inspector(SAGEConv)
+    assert str(inspector) == 'Inspector(SAGEConv)'
+    assert inspector.implements('message')
+    assert inspector.implements('message_and_aggregate')
+
+    out = inspector.inspect_signature(SAGEConv.message)
+    assert isinstance(out, Signature)
+    assert out.param_dict == {'x_j': Parameter('x_j', Tensor, inspect._empty)}
+    assert out.return_type == Tensor
+    assert inspector.get_flat_params(['message', 'message']) == [
+        Parameter('x_j', Tensor, inspect._empty),
+    ]
+    assert inspector.get_flat_param_names(['message']) == ['x_j']
+
+    kwargs = {'x_j': torch.randn(5), 'x_i': torch.randn(5)}
+    data = inspector.collect_param_data('message', kwargs)
+    assert len(data) == 1
+    assert torch.allclose(data['x_j'], kwargs['x_j'])
+
+    assert inspector.get_params_from_method_call(SAGEConv.propagate) == {
+        'x': Parameter('x', 'OptPairTensor', inspect._empty),
+    }
+
+
+def test_inspector_gat_conv() -> None:
+    inspector = Inspector(GATConv)
+    assert str(inspector) == 'Inspector(GATConv)'
+    assert inspector.implements('message')
+    assert not inspector.implements('message_and_aggregate')
+
+    out = inspector.inspect_signature(GATConv.message)
+    assert isinstance(out, Signature)
+    assert out.param_dict == {
+        'x_j': Parameter('x_j', Tensor, inspect._empty),
+        'alpha': Parameter('alpha', Tensor, inspect._empty),
+    }
+    assert out.return_type == Tensor
+    assert inspector.get_flat_params(['message', 'message']) == [
+        Parameter('x_j', Tensor, inspect._empty),
+        Parameter('alpha', Tensor, inspect._empty),
+    ]
+    assert inspector.get_flat_param_names(['message']) == ['x_j', 'alpha']
+
+    kwargs = {'x_j': torch.randn(5), 'alpha': torch.randn(5)}
+    data = inspector.collect_param_data('message', kwargs)
+    assert len(data) == 2
+    assert torch.allclose(data['x_j'], kwargs['x_j'])
+    assert torch.allclose(data['alpha'], kwargs['alpha'])
+
+    assert inspector.get_params_from_method_call(SAGEConv.propagate) == {
+        'x': Parameter('x', 'OptPairTensor', inspect._empty),
+        'alpha': Parameter('alpha', 'Tensor', inspect._empty),
+    }
+
+
+def test_get_params_from_method_call() -> None:
+    class FromMethodCall1:
+        propagate_type = {'x': Tensor}
+
+    inspector = Inspector(FromMethodCall1)
+    assert inspector.get_params_from_method_call('propagate') == {
+        'x': Parameter('x', Tensor, inspect._empty),
+    }
+
+    class FromMethodCall2:
+        # propagate_type: (x: Tensor)
+        pass
+
+    inspector = Inspector(FromMethodCall2)
+    assert inspector.get_params_from_method_call('propagate') == {
+        'x': Parameter('x', 'Tensor', inspect._empty),
+    }
+
+    class FromMethodCall3:
+        def forward(self) -> None:
+            self.propagate(  # type: ignore
+                torch.randn(5, 5),
+                x=None,
+                size=None,
+            )
+
+    inspector = Inspector(FromMethodCall3)
+    exclude = [0, 'size']
+    assert inspector.get_params_from_method_call('propagate', exclude) == {
+        'x': Parameter('x', Tensor, inspect._empty),
+    }
+
+    class FromMethodCall4:
+        pass
+
+    inspector = Inspector(FromMethodCall4)
+    assert inspector.get_params_from_method_call('propagate') == {}