Repository: pyg-team/pytorch_geometric
Pull request #9460 (2024-06-26 06:54:55)

Add `ConfigMixin`


Patch:
diff --git a/torch_geometric/config_mixin.py b/torch_geometric/config_mixin.py
new file mode 100644
index 000000000..6dbe5c430
--- /dev/null
+++ b/torch_geometric/config_mixin.py
@@ -0,0 +1,110 @@
+import inspect
+from dataclasses import fields, is_dataclass
+from importlib import import_module
+from typing import Any, Dict
+
+from torch_geometric.config_store import (
+    class_from_dataclass,
+    dataclass_from_class,
+)
+from torch_geometric.isinstance import is_torch_instance
+
+
+class ConfigMixin:
+    r"""Enables a class to serialize/deserialize itself to a dataclass."""
+    def config(self) -> Any:
+        r"""Creates a serializable configuration of the class."""
+        data_cls = dataclass_from_class(self.__class__)
+        if data_cls is None:
+            raise ValueError(f"Could not find the configuration class that "
+                             f"belongs to '{self.__class__.__name__}'. Please "
+                             f"register it in the configuration store.")
+
+        kwargs: Dict[str, Any] = {}
+        for field in fields(data_cls):
+            if not hasattr(self, field.name):
+                continue
+            kwargs[field.name] = _recursive_config(getattr(self, field.name))
+        return data_cls(**kwargs)
+
+    @classmethod
+    def from_config(cls, cfg: Any, *args: Any, **kwargs: Any) -> Any:
+        r"""Instantiates the class from a serializable configuration."""
+        if getattr(cfg, '_target_', None):
+            cls = _locate_cls(cfg._target_)
+        elif isinstance(cfg, dict) and '_target_' in cfg:
+            cls = _locate_cls(cfg['target'])
+
+        data_cls = cfg.__class__
+        if not is_dataclass(data_cls):
+            data_cls = dataclass_from_class(cls)
+            if data_cls is None:
+                raise ValueError(f"Could not find the configuration class "
+                                 f"that belongs to '{cls.__name__}'. Please "
+                                 f"register it in the configuration store.")
+
+        field_names = {field.name for field in fields(data_cls)}
+        if isinstance(cfg, dict):
+            _kwargs = {k: v for k, v in cfg.items() if k in field_names}
+            cfg = data_cls(**_kwargs)
+        assert is_dataclass(cfg)
+
+        if len(args) > 0:  # Convert `*args` to `**kwargs`:
+            param_names = list(inspect.signature(cls).parameters.keys())
+            if 'args' in param_names:
+                param_names.remove('args')
+            if 'kwargs' in param_names:
+                param_names.remove('kwargs')
+
+            for name, arg in zip(param_names, args):
+                kwargs[name] = arg
+
+        for key in field_names:
+            if key not in kwargs and key != '_target_':
+                kwargs[key] = _recursive_from_config(getattr(cfg, key))
+
+        return cls(**kwargs)
+
+
+def _recursive_config(value: Any) -> Any:
+    if isinstance(value, ConfigMixin):
+        return value.config()
+    if is_torch_instance(value, ConfigMixin):
+        return value.config()
+    if isinstance(value, (tuple, list)):
+        return [_recursive_config(v) for v in value]
+    if isinstance(value, dict):
+        return {k: _recursive_config(v) for k, v in value.items()}
+    return value
+
+
+def _recursive_from_config(value: Any) -> Any:
+    cls: Any = None
+    if is_dataclass(value):
+        if getattr(value, '_target_', None):
+            cls = _locate_cls(value._target_)
+        else:
+            cls = class_from_dataclass(value.__class__)
+    elif isinstance(value, dict) and '_target_' in value:
+        cls = _locate_cls(value['_target_'])
+
+    if cls is not None and issubclass(cls, ConfigMixin):
+        return cls.from_config(value)
+    if isinstance(value, (tuple, list)):
+        return [_recursive_from_config(v) for v in value]
+    if isinstance(value, dict):
+        return {k: _recursive_from_config(v) for k, v in value.items()}
+    return value
+
+
+def _locate_cls(qualname: str) -> Any:
+    parts = qualname.split('.')
+
+    if len(parts) <= 1:
+        raise ValueError(f"Qualified name is missing a dot (got '{qualname}')")
+
+    if any([len(part) == 0 for part in parts]):
+        raise ValueError(f"Relative imports not supported (got '{qualname}')")
+
+    module_name, cls_name = '.'.join(parts[:-1]), parts[-1]
+    return getattr(import_module(module_name), cls_name)


Test patch:
diff --git a/test/test_config_mixin.py b/test/test_config_mixin.py
new file mode 100644
index 000000000..fbff499fd
--- /dev/null
+++ b/test/test_config_mixin.py
@@ -0,0 +1,70 @@
+from dataclasses import dataclass, is_dataclass
+
+import torch
+
+from torch_geometric.config_mixin import ConfigMixin
+from torch_geometric.config_store import clear_config_store, register
+
+
+def teardown_function() -> None:
+    clear_config_store()
+
+
+@dataclass
+class Dataclass:
+    x: int
+    y: int
+
+
+class Base(torch.nn.Module, ConfigMixin):
+    pass
+
+
+@register(with_target=True)
+class Module(Base):
+    def __init__(self, x: int, data: Dataclass):
+        super().__init__()
+        self.x = x
+        self.data = data
+
+
+def test_config_mixin() -> None:
+    x = 0
+    data = Dataclass(x=1, y=2)
+
+    model = Module(x, data)
+    cfg = model.config()
+    assert is_dataclass(cfg)
+    assert cfg.x == 0
+    assert isinstance(cfg.data, Dataclass)
+    assert cfg.data.x == 1
+    assert cfg.data.y == 2
+    assert cfg._target_ == 'test_config_mixin.Module'
+
+    model = Module.from_config(cfg)
+    assert isinstance(model, Module)
+    assert model.x == 0
+    assert isinstance(model.data, Dataclass)
+    assert model.data.x == 1
+    assert model.data.y == 2
+
+    model = Base.from_config(cfg)
+    assert isinstance(model, Module)
+    assert model.x == 0
+    assert isinstance(model.data, Dataclass)
+    assert model.data.x == 1
+    assert model.data.y == 2
+
+    model = Base.from_config(cfg, 3)
+    assert isinstance(model, Module)
+    assert model.x == 3
+    assert isinstance(model.data, Dataclass)
+    assert model.data.x == 1
+    assert model.data.y == 2
+
+    model = Base.from_config(cfg, data=Dataclass(x=2, y=3))
+    assert isinstance(model, Module)
+    assert model.x == 0
+    assert isinstance(model.data, Dataclass)
+    assert model.data.x == 2
+    assert model.data.y == 3