Improve pytest detection using PYTEST_CURRENT_TEST

.github/workflows/documentation.yml

@@ -29,7 +29,6 @@ jobs:
         with:
           files: |
             benchmark/**
-            conda/**
             docker/**
             examples/**
             graphgym/**

.github/workflows/linting.yml

@@ -29,7 +29,6 @@ jobs:
         with:
           files: |
             benchmark/**
-            conda/**
             docker/**
             docs/**
             examples/**

.github/workflows/testing_latest.yml

@@ -41,7 +41,6 @@ jobs:
         with:
           files: |
             benchmark/**
-            conda/**
             docker/**
             docs/**
             examples/**

.github/workflows/testing_minimal.yml

@@ -29,7 +29,6 @@ jobs:
         with:
           files: |
             benchmark/**
-            conda/**
             docker/**
             docs/**
             examples/**

.github/workflows/testing_nightly.yml

@@ -29,7 +29,6 @@ jobs:
         with:
           files: |
             benchmark/**
-            conda/**
             docker/**
             docs/**
             examples/**

.github/workflows/testing_prev.yml

@@ -41,7 +41,6 @@ jobs:
         with:
           files: |
             benchmark/**
-            conda/**
             docker/**
             docs/**
             examples/**

.pre-commit-config.yaml

@@ -19,13 +19,6 @@ repos:
         name: Remove trailing whitespaces
       - id: check-toml
         name: Check toml
-      - id: check-yaml
-        name: Check yaml
-        exclude: |
-          (?x)^(
-              conda/pytorch-geometric/meta.yaml|
-              conda/pyg/meta.yaml
-          )$
 
   - repo: https://github.com/adrienverge/yamllint.git
     rev: v1.37.1

torch_geometric/data/dataset.py

@@ -1,4 +1,5 @@
 import copy
+import os
 import os.path as osp
 import re
 import sys
@@ -258,7 +259,7 @@ def _process(self):
         if not self.force_reload and files_exist(self.processed_paths):
             return
 
-        if self.log and 'pytest' not in sys.modules:
+        if self.log and 'PYTEST_CURRENT_TEST' not in os.environ:
             print('Processing...', file=sys.stderr)
 
         fs.makedirs(self.processed_dir, exist_ok=True)
@@ -269,7 +270,7 @@ def _process(self):
         path = osp.join(self.processed_dir, 'pre_filter.pt')
         fs.torch_save(_repr(self.pre_filter), path)
 
-        if self.log and 'pytest' not in sys.modules:
+        if self.log and 'PYTEST_CURRENT_TEST' not in os.environ:
             print('Done!', file=sys.stderr)
 
     def __len__(self) -> int:

torch_geometric/data/download.py

@@ -34,11 +34,11 @@ def download_url(
     path = osp.join(folder, filename)
 
     if fs.exists(path):  # pragma: no cover
-        if log and 'pytest' not in sys.modules:
+        if log and 'PYTEST_CURRENT_TEST' not in os.environ:
             print(f'Using existing file {filename}', file=sys.stderr)
         return path
 
-    if log and 'pytest' not in sys.modules:
+    if log and 'PYTEST_CURRENT_TEST' not in os.environ:
         print(f'Downloading {url}', file=sys.stderr)
 
     os.makedirs(folder, exist_ok=True)

torch_geometric/data/extract.py

@@ -1,13 +1,14 @@
 import bz2
 import gzip
+import os
 import os.path as osp
 import sys
 import tarfile
 import zipfile
 
 
 def maybe_log(path: str, log: bool = True) -> None:
-    if log and 'pytest' not in sys.modules:
+    if log and 'PYTEST_CURRENT_TEST' not in os.environ:
         print(f'Extracting {path}', file=sys.stderr)
 
 

torch_geometric/io/fs.py

@@ -1,4 +1,5 @@
 import io
+import os
 import os.path as osp
 import pickle
 import re
@@ -117,7 +118,7 @@ def cp(
     # Cache result if the protocol is not local:
     cache_dir: Optional[str] = None
     if not islocal(path1):
-        if log and 'pytest' not in sys.modules:
+        if log and 'PYTEST_CURRENT_TEST' not in os.environ:
             print(f'Downloading {path1}', file=sys.stderr)
 
         if extract and use_cache:  # Cache seems to confuse the gcs filesystem.

torch_geometric/nn/conv/fa_conv.py

@@ -163,10 +163,13 @@ def forward(  # noqa: F811
             edge_index (torch.Tensor or SparseTensor): The edge indices.
             edge_weight (torch.Tensor, optional): The edge weights.
                 (default: :obj:`None`)
-            return_attention_weights (bool, optional): If set to :obj:`True`,
-                will additionally return the tuple
-                :obj:`(edge_index, attention_weights)`, holding the computed
-                attention weights for each edge. (default: :obj:`None`)
+            return_attention_weights (bool, optional):
+                Will additionally return the tuple
+                :obj:`(edge_index, attention_weights)` whenever it is set to
+                a value, regardless of its actual value
+                (might be `True` or `False`), holding the computed attention
+                weights for each edge.
+                (default: :obj:`None`)
         """
         if self.normalize:
             if isinstance(edge_index, Tensor):

torch_geometric/nn/conv/gat_conv.py

@@ -273,10 +273,13 @@ def forward(  # noqa: F811
                 (default: :obj:`None`)
             size ((int, int), optional): The shape of the adjacency matrix.
                 (default: :obj:`None`)
-            return_attention_weights (bool, optional): If set to :obj:`True`,
-                will additionally return the tuple
-                :obj:`(edge_index, attention_weights)`, holding the computed
-                attention weights for each edge. (default: :obj:`None`)
+            return_attention_weights (bool, optional):
+                Will additionally return the tuple
+                :obj:`(edge_index, attention_weights)` whenever it is set to
+                a value, regardless of its actual value
+                (might be `True` or `False`), holding the computed attention
+                weights for each edge.
+                (default: :obj:`None`)
         """
         # NOTE: attention weights will be returned whenever
         # `return_attention_weights` is set to a value, regardless of its

torch_geometric/nn/conv/gatv2_conv.py

@@ -266,10 +266,13 @@ def forward(  # noqa: F811
             edge_index (torch.Tensor or SparseTensor): The edge indices.
             edge_attr (torch.Tensor, optional): The edge features.
                 (default: :obj:`None`)
-            return_attention_weights (bool, optional): If set to :obj:`True`,
-                will additionally return the tuple
-                :obj:`(edge_index, attention_weights)`, holding the computed
-                attention weights for each edge. (default: :obj:`None`)
+            return_attention_weights (bool, optional):
+                Will additionally return the tuple
+                :obj:`(edge_index, attention_weights)` whenever it is set to
+                a value, regardless of its actual value
+                (might be `True` or `False`), holding the computed attention
+                weights for each edge.
+                (default: :obj:`None`)
         """
         H, C = self.heads, self.out_channels
 

torch_geometric/nn/conv/rgat_conv.py

@@ -340,10 +340,13 @@ def forward(
                 (default: :obj:`None`)
             size ((int, int), optional): The shape of the adjacency matrix.
                 (default: :obj:`None`)
-            return_attention_weights (bool, optional): If set to :obj:`True`,
-                will additionally return the tuple
-                :obj:`(edge_index, attention_weights)`, holding the computed
-                attention weights for each edge. (default: :obj:`None`)
+            return_attention_weights (bool, optional):
+                Will additionally return the tuple
+                :obj:`(edge_index, attention_weights)` whenever it is set to
+                a value, regardless of its actual value
+                (might be `True` or `False`), holding the computed attention
+                weights for each edge.
+                (default: :obj:`None`)
         """
         # propagate_type: (x: Tensor, edge_type: OptTensor,
         #                  edge_attr: OptTensor)

torch_geometric/nn/conv/transformer_conv.py

@@ -212,10 +212,13 @@ def forward(  # noqa: F811
             edge_index (torch.Tensor or SparseTensor): The edge indices.
             edge_attr (torch.Tensor, optional): The edge features.
                 (default: :obj:`None`)
-            return_attention_weights (bool, optional): If set to :obj:`True`,
-                will additionally return the tuple
-                :obj:`(edge_index, attention_weights)`, holding the computed
-                attention weights for each edge. (default: :obj:`None`)
+            return_attention_weights (bool, optional):
+                Will additionally return the tuple
+                :obj:`(edge_index, attention_weights)` whenever it is set to
+                a value, regardless of its actual value
+                (might be `True` or `False`), holding the computed attention
+                weights for each edge.
+                (default: :obj:`None`)
         """
         H, C = self.heads, self.out_channels
 

torch_geometric/nn/dense/linear.py

@@ -1,5 +1,5 @@
 import math
-import sys
+import os
 import time
 from typing import Any, Dict, Optional, Tuple, Union
 
@@ -262,7 +262,7 @@ def _update_timing_cache(
         key: int,
     ) -> None:
 
-        MEASURE_ITER = 1 if 'pytest' in sys.modules else 3
+        MEASURE_ITER = 1 if 'PYTEST_CURRENT_TEST' not in os.environ else 3
 
         if torch.cuda.is_available():
             torch.cuda.synchronize()

torch_geometric/typing.py

@@ -1,6 +1,5 @@
 import inspect
 import os
-import sys
 import typing
 import warnings
 from typing import Any, Dict, List, Optional, Set, Tuple, TypeAlias, Union
@@ -43,7 +42,8 @@
     WITH_PYG_LIB = True
     WITH_GMM = WITH_PT20 and hasattr(pyg_lib.ops, 'grouped_matmul')
     WITH_SEGMM = hasattr(pyg_lib.ops, 'segment_matmul')
-    if WITH_SEGMM and 'pytest' in sys.modules and torch.cuda.is_available():
+    if WITH_SEGMM and torch.cuda.is_available(
+    ) and 'PYTEST_CURRENT_TEST' not in os.environ:
         # NOTE `segment_matmul` is currently bugged on older NVIDIA cards which
         # let our GPU tests on CI crash. Try if this error is present on the
         # current GPU and disable `WITH_SEGMM`/`WITH_GMM` if necessary.