Repository: pyg-team/pytorch_geometric
Pull request #8614 (2023-12-13 20:50:49)

Support `mypy` in `torch_geometric.transforms` [1/n]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8614?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:
> Comparison is base [(`de8f561`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/de8f5615b789effb357c687b087c233eb96faec5?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.56% compared to head [(`eaece73`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8614?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.89%.

> :exclamation: Current head eaece73 differs from pull request most recent head f549cda. Consider uploading reports for the commit f549cda to get more accurate results


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


```diff
@@            Coverage Diff             @@
##           master    #8614      +/-   ##
==========================================
- Coverage   89.56%   88.89%   -0.67%     
==========================================
  Files         480      480              
  Lines       30532    30549      +17     
==========================================
- Hits        27345    27158     -187     
- Misses       3187     3391     +204     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8614?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/torch_geometric/data/data.py b/torch_geometric/data/data.py
index 97a4c8d2a..2d71bb46f 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -999,6 +999,14 @@ class Data(BaseData, FeatureStore, GraphStore):
     def time(self, time: Optional[Tensor]):
         self._store.time = time
 
+    @property
+    def face(self) -> Optional[Tensor]:
+        return self['face'] if 'face' in self._store else None
+
+    @face.setter
+    def face(self, face: Optional[Tensor]):
+        self._store.face = face
+
     # Deprecated functions ####################################################
 
     @property
diff --git a/torch_geometric/data/storage.py b/torch_geometric/data/storage.py
index 6f6b0a42c..1bdaaf177 100644
--- a/torch_geometric/data/storage.py
+++ b/torch_geometric/data/storage.py
@@ -15,6 +15,7 @@ from typing import (
     Set,
     Tuple,
     Union,
+    overload,
 )
 
 import numpy as np
@@ -538,6 +539,10 @@ class EdgeStorage(BaseStorage):
             f"'{self.__class__.__name__}' object has no attribute "
             f"'edge_index', 'adj' or 'adj_t'")
 
+    @edge_index.setter
+    def edge_index(self, edge_index: Optional[Tensor]) -> None:
+        self['edge_index'] = edge_index
+
     @property
     def num_edges(self) -> int:
         # We sequentially access attributes that reveal the number of edges.
@@ -582,6 +587,14 @@ class EdgeStorage(BaseStorage):
     def num_features(self) -> int:
         return self.num_edge_features
 
+    @overload
+    def size(self) -> Tuple[Optional[int], Optional[int]]:
+        pass
+
+    @overload
+    def size(self, dim: int) -> Optional[int]:
+        pass
+
     def size(
         self, dim: Optional[int] = None
     ) -> Union[Tuple[Optional[int], Optional[int]], Optional[int]]:
diff --git a/torch_geometric/transforms/add_metapaths.py b/torch_geometric/transforms/add_metapaths.py
index 427a62a42..94f295489 100644
--- a/torch_geometric/transforms/add_metapaths.py
+++ b/torch_geometric/transforms/add_metapaths.py
@@ -104,9 +104,8 @@ class AddMetaPaths(BaseTransform):
         drop_unconnected_node_types: bool = False,
         max_sample: Optional[int] = None,
         weighted: bool = False,
-        **kwargs,
-    ):
-
+        **kwargs: bool,
+    ) -> None:
         if 'drop_orig_edges' in kwargs:
             warnings.warn("'drop_orig_edges' is dprecated. Use "
                           "'drop_orig_edge_types' instead")
@@ -298,9 +297,13 @@ class AddRandomMetaPaths(BaseTransform):
                 f'walks_per_node={self.walks_per_node})')
 
 
-def postprocess(data: HeteroData, edge_types: List[EdgeType],
-                drop_orig_edge_types: bool, keep_same_node_type: bool,
-                drop_unconnected_node_types: bool):
+def postprocess(
+    data: HeteroData,
+    edge_types: List[EdgeType],
+    drop_orig_edge_types: bool,
+    keep_same_node_type: bool,
+    drop_unconnected_node_types: bool,
+) -> None:
 
     if drop_orig_edge_types:
         for i in edge_types:
diff --git a/torch_geometric/transforms/add_remaining_self_loops.py b/torch_geometric/transforms/add_remaining_self_loops.py
index 4150f3025..ec3de0cda 100644
--- a/torch_geometric/transforms/add_remaining_self_loops.py
+++ b/torch_geometric/transforms/add_remaining_self_loops.py
@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Union
 
 from torch import Tensor
 
@@ -27,8 +27,11 @@ class AddRemainingSelfLoops(BaseTransform):
             according to a reduce operation. (:obj:`"add"`, :obj:`"mean"`,
             :obj:`"min"`, :obj:`"max"`, :obj:`"mul"`). (default: :obj:`1.`)
     """
-    def __init__(self, attr: Optional[str] = 'edge_weight',
-                 fill_value: Union[float, Tensor, str] = 1.0):
+    def __init__(
+        self,
+        attr: str = 'edge_weight',
+        fill_value: Union[float, Tensor, str] = 1.0,
+    ):
         self.attr = attr
         self.fill_value = fill_value
 
@@ -40,10 +43,11 @@ class AddRemainingSelfLoops(BaseTransform):
             if store.is_bipartite() or 'edge_index' not in store:
                 continue
 
-            store.edge_index, edge_weight = add_remaining_self_loops(
-                store.edge_index, getattr(store, self.attr, None),
-                fill_value=self.fill_value, num_nodes=store.size(0))
-
-            setattr(store, self.attr, edge_weight)
+            store.edge_index, store[self.attr] = add_remaining_self_loops(
+                store.edge_index,
+                edge_attr=store.get(self.attr, None),
+                fill_value=self.fill_value,
+                num_nodes=store.size(0),
+            )
 
         return data
diff --git a/torch_geometric/transforms/add_self_loops.py b/torch_geometric/transforms/add_self_loops.py
index c6594970a..7fe871377 100644
--- a/torch_geometric/transforms/add_self_loops.py
+++ b/torch_geometric/transforms/add_self_loops.py
@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Union
 
 from torch import Tensor
 
@@ -27,8 +27,11 @@ class AddSelfLoops(BaseTransform):
             according to a reduce operation. (:obj:`"add"`, :obj:`"mean"`,
             :obj:`"min"`, :obj:`"max"`, :obj:`"mul"`). (default: :obj:`1.`)
     """
-    def __init__(self, attr: Optional[str] = 'edge_weight',
-                 fill_value: Union[float, Tensor, str] = 1.0):
+    def __init__(
+        self,
+        attr: str = 'edge_weight',
+        fill_value: Union[float, Tensor, str] = 1.0,
+    ) -> None:
         self.attr = attr
         self.fill_value = fill_value
 
@@ -40,10 +43,11 @@ class AddSelfLoops(BaseTransform):
             if store.is_bipartite() or 'edge_index' not in store:
                 continue
 
-            store.edge_index, edge_weight = add_self_loops(
-                store.edge_index, getattr(store, self.attr, None),
-                fill_value=self.fill_value, num_nodes=store.size(0))
-
-            setattr(store, self.attr, edge_weight)
+            store.edge_index, store[self.attr] = add_self_loops(
+                store.edge_index,
+                edge_attr=store.get(self.attr, None),
+                fill_value=self.fill_value,
+                num_nodes=store.size(0),
+            )
 
         return data
diff --git a/torch_geometric/transforms/cartesian.py b/torch_geometric/transforms/cartesian.py
index ec5fd5dea..329d308f7 100644
--- a/torch_geometric/transforms/cartesian.py
+++ b/torch_geometric/transforms/cartesian.py
@@ -37,17 +37,19 @@ class Cartesian(BaseTransform):
         self.interval = interval
 
     def forward(self, data: Data) -> Data:
+        assert data.pos is not None
+        assert data.edge_index is not None
         (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
 
         cart = pos[row] - pos[col]
         cart = cart.view(-1, 1) if cart.dim() == 1 else cart
 
         if self.norm and cart.numel() > 0:
-            max_value = cart.abs().max() if self.max is None else self.max
+            max_val = float(cart.abs().max()) if self.max is None else self.max
 
             length = self.interval[1] - self.interval[0]
             center = (self.interval[0] + self.interval[1]) / 2
-            cart = length * cart / (2 * max_value) + center
+            cart = length * cart / (2 * max_val) + center
 
         if pseudo is not None and self.cat:
             pseudo = pseudo.view(-1, 1) if pseudo.dim() == 1 else pseudo
diff --git a/torch_geometric/transforms/constant.py b/torch_geometric/transforms/constant.py
index 2cfe87df7..da172f61a 100644
--- a/torch_geometric/transforms/constant.py
+++ b/torch_geometric/transforms/constant.py
@@ -42,6 +42,7 @@ class Constant(BaseTransform):
         for store in data.node_stores:
             if self.node_types is None or store._key in self.node_types:
                 num_nodes = store.num_nodes
+                assert num_nodes is not None
                 c = torch.full((num_nodes, 1), self.value, dtype=torch.float)
 
                 if hasattr(store, 'x') and self.cat:
diff --git a/torch_geometric/transforms/distance.py b/torch_geometric/transforms/distance.py
index d012c4d2f..8d6692e81 100644
--- a/torch_geometric/transforms/distance.py
+++ b/torch_geometric/transforms/distance.py
@@ -37,15 +37,17 @@ class Distance(BaseTransform):
         self.interval = interval
 
     def forward(self, data: Data) -> Data:
+        assert data.pos is not None
+        assert data.edge_index is not None
         (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
 
         dist = torch.norm(pos[col] - pos[row], p=2, dim=-1).view(-1, 1)
 
         if self.norm and dist.numel() > 0:
-            max_value = dist.max() if self.max is None else self.max
+            max_val = float(dist.max()) if self.max is None else self.max
 
             length = self.interval[1] - self.interval[0]
-            dist = length * (dist / max_value) + self.interval[0]
+            dist = length * (dist / max_val) + self.interval[0]
 
         if pseudo is not None and self.cat:
             pseudo = pseudo.view(-1, 1) if pseudo.dim() == 1 else pseudo
diff --git a/torch_geometric/transforms/fixed_points.py b/torch_geometric/transforms/fixed_points.py
index 0340dd835..7197fde79 100644
--- a/torch_geometric/transforms/fixed_points.py
+++ b/torch_geometric/transforms/fixed_points.py
@@ -3,6 +3,7 @@ import re
 
 import numpy as np
 import torch
+from torch import Tensor
 
 from torch_geometric.data import Data
 from torch_geometric.data.datapipes import functional_transform
@@ -39,10 +40,11 @@ class FixedPoints(BaseTransform):
 
     def forward(self, data: Data) -> Data:
         num_nodes = data.num_nodes
+        assert num_nodes is not None
 
         if self.replace:
-            choice = np.random.choice(num_nodes, self.num, replace=True)
-            choice = torch.from_numpy(choice).to(torch.long)
+            choice = torch.from_numpy(
+                np.random.choice(num_nodes, self.num, replace=True)).long()
         elif not self.allow_duplicates:
             choice = torch.randperm(num_nodes)[:self.num]
         else:
@@ -51,14 +53,14 @@ class FixedPoints(BaseTransform):
                 for _ in range(math.ceil(self.num / num_nodes))
             ], dim=0)[:self.num]
 
-        for key, item in data:
+        for key, value in data.items():
             if key == 'num_nodes':
                 data.num_nodes = choice.size(0)
             elif bool(re.search('edge', key)):
                 continue
-            elif (torch.is_tensor(item) and item.size(0) == num_nodes
-                  and item.size(0) != 1):
-                data[key] = item[choice]
+            elif (isinstance(value, Tensor) and value.size(0) == num_nodes
+                  and value.size(0) != 1):
+                data[key] = value[choice]
 
         return data
 
diff --git a/torch_geometric/transforms/generate_mesh_normals.py b/torch_geometric/transforms/generate_mesh_normals.py
index 5440ee82e..f5def792f 100644
--- a/torch_geometric/transforms/generate_mesh_normals.py
+++ b/torch_geometric/transforms/generate_mesh_normals.py
@@ -12,7 +12,8 @@ class GenerateMeshNormals(BaseTransform):
     faces (functional name: :obj:`generate_mesh_normals`).
     """
     def forward(self, data: Data) -> Data:
-        assert 'face' in data
+        assert data.pos is not None
+        assert data.face is not None
         pos, face = data.pos, data.face
 
         vec1 = pos[face[1]] - pos[face[0]]
diff --git a/torch_geometric/transforms/local_cartesian.py b/torch_geometric/transforms/local_cartesian.py
index d36455ed5..cb00e444c 100644
--- a/torch_geometric/transforms/local_cartesian.py
+++ b/torch_geometric/transforms/local_cartesian.py
@@ -34,6 +34,8 @@ class LocalCartesian(BaseTransform):
         self.interval = interval
 
     def forward(self, data: Data) -> Data:
+        assert data.pos is not None
+        assert data.edge_index is not None
         (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
 
         cart = pos[row] - pos[col]
diff --git a/torch_geometric/transforms/spherical.py b/torch_geometric/transforms/spherical.py
index 9f68637ae..182aba1be 100644
--- a/torch_geometric/transforms/spherical.py
+++ b/torch_geometric/transforms/spherical.py
@@ -34,6 +34,8 @@ class Spherical(BaseTransform):
         self.cat = cat
 
     def forward(self, data: Data) -> Data:
+        assert data.pos is not None
+        assert data.edge_index is not None
         (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
         assert pos.dim() == 2 and pos.size(1) == 3
 


Test patch:
diff --git a/test/test_config_store.py b/test/test_config_store.py
index cb753a0b2..37e5ea4cc 100644
--- a/test/test_config_store.py
+++ b/test/test_config_store.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 from torch_geometric.config_store import (
     class_from_dataclass,
@@ -26,7 +26,7 @@ def test_to_dataclass():
     fields = AddSelfLoopsConfig.__dataclass_fields__
 
     assert fields['attr'].name == 'attr'
-    assert fields['attr'].type == Optional[str]
+    assert fields['attr'].type == str
     assert fields['attr'].default == 'edge_weight'
 
     assert fields['fill_value'].name == 'fill_value'