Repository: pyg-team/pytorch_geometric
Pull request #9269 (2024-05-02 09:18:03)

Fix `map_annotation` when annotation is `{list,dict}` instead of `typing.{List,Dict}`
Fixes this error:

```python
from dataclasses import dataclass
from typing import Dict

from torch_geometric.config_store import register


@register()
@dataclass
class MyClass1:
    something: Dict[str, int]


@register()
@dataclass
class MyClass2:
    something: dict[str, int]
```

```
Traceback (most recent call last):
  File "/home/aki/work/github.com/pyg-team/pytorch_geometric/test_config.py", line 13, in <module>
    @register()
     ^^^^^^^^^^
  File "/home/aki/work/github.com/pyg-team/pytorch_geometric/torch_geometric/config_store.py", line 360, in bounded_register
    register(cls=cls, data_cls=data_cls, group=group, **kwargs)
  File "/home/aki/work/github.com/pyg-team/pytorch_geometric/torch_geometric/config_store.py", line 345, in register
    data_cls = to_dataclass(cls, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/aki/work/github.com/pyg-team/pytorch_geometric/torch_geometric/config_store.py", line 255, in to_dataclass
    annotation = map_annotation(annotation, mapping=MAPPING)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/aki/work/github.com/pyg-team/pytorch_geometric/torch_geometric/config_store.py", line 170, in map_annotation
    annotation.__args__ = tuple(map_annotation(a, mapping) for a in args)
    ^^^^^^^^^^^^^^^^^^^
AttributeError: readonly attribute
```

akihironitta: TODO: handle the case with older Python versions

Patch:
diff --git a/torch_geometric/config_store.py b/torch_geometric/config_store.py
index 6b746a124..4a54b23f6 100644
--- a/torch_geometric/config_store.py
+++ b/torch_geometric/config_store.py
@@ -162,12 +162,18 @@ def map_annotation(
     annotation: Any,
     mapping: Optional[Dict[Any, Any]] = None,
 ) -> Any:
-
     origin = getattr(annotation, '__origin__', None)
-    args = getattr(annotation, '__args__', [])
-    if origin == Union or origin == list or origin == dict:
-        annotation = copy.copy(annotation)
-        annotation.__args__ = tuple(map_annotation(a, mapping) for a in args)
+    args = getattr(annotation, '__args__', tuple())
+    if origin in {Union, list, dict, tuple}:
+        new_args = tuple(map_annotation(a, mapping) for a in args)
+        if type(annotation).__name__ == 'GenericAlias':
+            # If annotated with `list[...]` or `dict[...]` (>= Python 3.10):
+            annotation = origin[new_args]
+        else:
+            # If annotated with `typing.List[...]` or `typing.Dict[...]`:
+            annotation = copy.copy(annotation)
+            annotation.__args__ = new_args
+
         return annotation
 
     if mapping is not None and annotation in mapping:


Test patch:
diff --git a/test/test_config_store.py b/test/test_config_store.py
index 37e5ea4cc..fcdb1196f 100644
--- a/test/test_config_store.py
+++ b/test/test_config_store.py
@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Dict, List, Tuple
 
 from torch_geometric.config_store import (
     class_from_dataclass,
@@ -6,10 +6,11 @@ from torch_geometric.config_store import (
     dataclass_from_class,
     fill_config_store,
     get_config_store,
+    map_annotation,
     register,
     to_dataclass,
 )
-from torch_geometric.testing import withPackage
+from torch_geometric.testing import minPython, withPackage
 from torch_geometric.transforms import AddSelfLoops
 
 
@@ -44,6 +45,21 @@ def test_to_dataclass():
                         "AddSelfLoops')")
 
 
+@minPython('3.10')
+def test_map_annotation():
+    mapping = {int: Any}
+    assert map_annotation(dict[str, int], mapping) == dict[str, Any]
+    assert map_annotation(Dict[str, float], mapping) == Dict[str, float]
+    assert map_annotation(List[str], mapping) == List[str]
+    assert map_annotation(List[int], mapping) == List[Any]
+    assert map_annotation(Tuple[int], mapping) == Tuple[Any]
+    assert map_annotation(dict[str, int], mapping) == dict[str, Any]
+    assert map_annotation(dict[str, float], mapping) == dict[str, float]
+    assert map_annotation(list[str], mapping) == list[str]
+    assert map_annotation(list[int], mapping) == list[Any]
+    assert map_annotation(tuple[int], mapping) == tuple[Any]
+
+
 def test_register():
     register(AddSelfLoops, group='transform')
     assert 'transform' in get_config_store().repo
diff --git a/torch_geometric/testing/__init__.py b/torch_geometric/testing/__init__.py
index 6da99347a..6f9b9baf2 100644
--- a/torch_geometric/testing/__init__.py
+++ b/torch_geometric/testing/__init__.py
@@ -10,7 +10,7 @@ from .decorators import (
     onlyDistributedTest,
     onlyLinux,
     noWindows,
-    onlyPython,
+    minPython,
     onlyCUDA,
     onlyXPU,
     onlyOnline,
@@ -40,7 +40,7 @@ __all__ = [
     'onlyDistributedTest',
     'onlyLinux',
     'noWindows',
-    'onlyPython',
+    'minPython',
     'onlyCUDA',
     'onlyXPU',
     'onlyOnline',
diff --git a/torch_geometric/testing/decorators.py b/torch_geometric/testing/decorators.py
index 5ecac30fa..3b197affb 100644
--- a/torch_geometric/testing/decorators.py
+++ b/torch_geometric/testing/decorators.py
@@ -67,15 +67,23 @@ def noWindows(func: Callable) -> Callable:
     )(func)
 
 
-def onlyPython(*args: str) -> Callable:
+def minPython(version: str) -> Callable:
     r"""A decorator to run tests on specific :python:`Python` versions only."""
     def decorator(func: Callable) -> Callable:
         import pytest
 
-        python_version = f'{sys.version_info.major}.{sys.version_info.minor}'
+        major, minor = version.split('.')
+
+        skip = False
+        if sys.version_info.major < int(major):
+            skip = True
+        if (sys.version_info.major == int(major)
+                and sys.version_info.minor < int(minor)):
+            skip = True
+
         return pytest.mark.skipif(
-            python_version not in args,
-            reason=f"Python {python_version} not supported",
+            skip,
+            reason=f"Python {version} required",
         )(func)
 
     return decorator