Repository: pyg-team/pytorch_geometric
Pull request #9044 (2024-03-11 07:57:37)

Ignore `self.propagate` appearances in comments when parsing `MessagePassing`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9044?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 89.26%. Comparing base [(`0e1b50b`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/0e1b50bdf3484769a1e1c80637984c7d047d90eb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) to head [(`54317e4`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9044?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).

> :exclamation: Current head 54317e4 differs from pull request most recent head 20a8681. Consider uploading reports for the commit 20a8681 to get more accurate results


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


```diff
@@            Coverage Diff             @@
##           master    #9044      +/-   ##
==========================================
- Coverage   89.88%   89.26%   -0.63%     
==========================================
  Files         468      468              
  Lines       30038    30041       +3     
==========================================
- Hits        26999    26815     -184     
- Misses       3039     3226     +187     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9044?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/CHANGELOG.md b/CHANGELOG.md
index a931f86de..e24273c36 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Ignore `self.propagate` appearances in comments when parsing `MessagePassing` implementation ([#9044](https://github.com/pyg-team/pytorch_geometric/pull/9044))
 - Fixed `OSError` on read-only file systems within `MessagePassing` ([#9032](https://github.com/pyg-team/pytorch_geometric/pull/9032))
 - Fixed metaclass conflict in `Dataset` ([#8999](https://github.com/pyg-team/pytorch_geometric/pull/8999))
 - Fixed import errors on `MessagePassing` modules with nested inheritance ([#8973](https://github.com/pyg-team/pytorch_geometric/pull/8973))
diff --git a/torch_geometric/inspector.py b/torch_geometric/inspector.py
index 1ebadd80a..e65b4d049 100644
--- a/torch_geometric/inspector.py
+++ b/torch_geometric/inspector.py
@@ -396,6 +396,7 @@ class Inspector:
         # (3) Parse the function call:
         for cls in self._cls.__mro__:
             source = self.get_source(cls)
+            source = remove_comments(source)
             match = find_parenthesis_content(source, f'self.{func_name}')
             if match is not None:
                 for i, kwarg in enumerate(split(match, sep=',')):
@@ -523,3 +524,12 @@ def split(content: str, sep: str) -> List[str]:
     if start != len(content):  # Respect dangling `sep`:
         outs.append(content[start:].strip())
     return outs
+
+
+def remove_comments(content: str) -> str:
+    content = re.sub(r'\s*#.*', '', content)
+    content = re.sub(re.compile(r'r"""(.*?)"""', re.DOTALL), '', content)
+    content = re.sub(re.compile(r'"""(.*?)"""', re.DOTALL), '', content)
+    content = re.sub(re.compile(r"r'''(.*?)'''", re.DOTALL), '', content)
+    content = re.sub(re.compile(r"'''(.*?)'''", re.DOTALL), '', content)
+    return content


Test patch:
diff --git a/test/nn/conv/test_message_passing.py b/test/nn/conv/test_message_passing.py
index 23812ebe9..901c7ac87 100644
--- a/test/nn/conv/test_message_passing.py
+++ b/test/nn/conv/test_message_passing.py
@@ -69,7 +69,7 @@ class MyConvWithSelfLoops(MessagePassing):
     def __init__(self, aggr: str = 'add'):
         super().__init__(aggr=aggr)
 
-    def forward(self, x: Tensor, edge_index: torch.Tensor) -> Tensor:
+    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
         edge_index, _ = add_self_loops(edge_index)
 
         # propagate_type: (x: Tensor)
@@ -141,6 +141,25 @@ def test_my_conv_edge_index():
     assert out.size() == (4, 32)
 
 
+class MyCommentedConv(MessagePassing):
+    r"""This layer calls `self.propagate()` internally."""
+    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
+        # `self.propagate()` is used here to propagate messages.
+        return self.propagate(edge_index, x=x)
+
+
+def test_my_commented_conv():
+    # Check that `self.propagate` occurences in comments are correctly ignored.
+    x = torch.randn(4, 8)
+    edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
+
+    conv = MyCommentedConv()
+    conv(x, edge_index)
+
+    jit = torch.jit.script(conv)
+    jit(x, edge_index)
+
+
 def test_my_conv_out_of_bounds():
     x = torch.randn(3, 8)
     value = torch.randn(4)