Backport dataclasses.asdict(…) for Python ≤ 3.11

Work around python/cpython#79721 on these versions.

Author: khaeru

message_ix_models/tests/util/test_context.py

@@ -13,6 +13,7 @@
 from message_ix import Scenario
 
 from message_ix_models import Context
+from message_ix_models.util.scenarioinfo import ScenarioInfo
 
 
 class TestContext:
@@ -155,12 +156,16 @@ def test_set_scenario(self, test_context):
             dict(model="foo", scenario="bar", version=0) == test_context.scenario_info
         )
 
-    def test_asdict(self, session_context):
+    def test_asdict(self, test_context):
+        # Add a ScenarioInfo object. This fails on Python <= 3.11 due to
+        # https://github.com/python/cpython/issues/79721
+        test_context.core.scenarios.append(ScenarioInfo())
+
         # asdict() method runs
-        session_context.asdict()
+        test_context.asdict()
 
         # Context can be serialized to json using the genno caching Encoder
-        json.dumps(session_context, cls=genno.caching.Encoder)
+        json.dumps(test_context, cls=genno.caching.Encoder)
 
     def test_write_debug_archive(self, mix_models_cli):
         """:meth:`.write_debug_archive` works."""

message_ix_models/tests/util/test_scenarioinfo.py

@@ -1,5 +1,8 @@
 import logging
 import re
+import sys
+from copy import deepcopy
+from dataclasses import asdict as asdict_stdlib
 
 import pandas as pd
 import pytest
@@ -12,9 +15,34 @@
 from message_ix_models import ScenarioInfo, Spec
 from message_ix_models.model.structure import get_codes, process_technology_codes
 from message_ix_models.util import as_codes
+from message_ix_models.util._dataclasses import asdict as asdict_backport
 
 
 class TestScenarioInfo:
+    @pytest.fixture(scope="class")
+    def info(self) -> ScenarioInfo:
+        return ScenarioInfo()
+
+    @pytest.mark.parametrize(
+        "func",
+        (
+            pytest.param(
+                asdict_stdlib,
+                marks=pytest.mark.xfail(
+                    condition=sys.version_info.minor <= 11,
+                    reason="https://github.com/python/cpython/issues/79721",
+                ),
+            ),
+            asdict_backport,
+        ),
+    )
+    def test_asdict(self, func, info) -> None:
+        """Test backported :func:`.asdict` works for ScenarioInfo."""
+        func(info)
+
+    def test_deepcopy(self, info) -> None:
+        deepcopy(info)
+
     def test_empty(self):
         """ScenarioInfo created from scratch."""
         info = ScenarioInfo()

message_ix_models/util/_dataclasses.py

@@ -0,0 +1,95 @@
+"""Utilities for :mod:`.dataclasses`.
+
+This module currently backports one function, :func:`asdict`, from Python 3.13.0, in
+order to avoid https://github.com/python/cpython/issues/79721. This issue specifically
+occurs with :attr:`.ScenarioInfo.set`, which is of class :class:`defaultdict`. The
+backported function **should** be used when (a) Python 3.11 or earlier is in use and (b)
+ScenarioInfo is handled directly or indirectly.
+"""
+# NB Comments are deleted
+
+import copy
+import types
+from dataclasses import fields
+
+__all__ = [
+    "asdict",
+]
+
+_ATOMIC_TYPES = frozenset(
+    {
+        types.NoneType,
+        bool,
+        int,
+        float,
+        str,
+        complex,
+        bytes,
+        types.EllipsisType,
+        types.NotImplementedType,
+        types.CodeType,
+        types.BuiltinFunctionType,
+        types.FunctionType,
+        type,
+        range,
+        property,
+    }
+)
+
+_FIELDS = "__dataclass_fields__"
+
+
+def _is_dataclass_instance(obj):
+    return hasattr(type(obj), _FIELDS)
+
+
+def asdict(obj, *, dict_factory=dict):
+    if not _is_dataclass_instance(obj):
+        raise TypeError("asdict() should be called on dataclass instances")
+    return _asdict_inner(obj, dict_factory)
+
+
+def _asdict_inner(obj, dict_factory):  # noqa: C901
+    obj_type = type(obj)
+    if obj_type in _ATOMIC_TYPES:
+        return obj
+    elif hasattr(obj_type, _FIELDS):
+        if dict_factory is dict:
+            return {
+                f.name: _asdict_inner(getattr(obj, f.name), dict) for f in fields(obj)
+            }
+        else:
+            return dict_factory(
+                [
+                    (f.name, _asdict_inner(getattr(obj, f.name), dict_factory))
+                    for f in fields(obj)
+                ]
+            )
+    elif obj_type is list:
+        return [_asdict_inner(v, dict_factory) for v in obj]
+    elif obj_type is dict:
+        return {
+            _asdict_inner(k, dict_factory): _asdict_inner(v, dict_factory)
+            for k, v in obj.items()
+        }
+    elif obj_type is tuple:
+        return tuple([_asdict_inner(v, dict_factory) for v in obj])
+    elif issubclass(obj_type, tuple):
+        if hasattr(obj, "_fields"):
+            return obj_type(*[_asdict_inner(v, dict_factory) for v in obj])
+        else:
+            return obj_type(_asdict_inner(v, dict_factory) for v in obj)
+    elif issubclass(obj_type, dict):
+        if hasattr(obj_type, "default_factory"):
+            result = obj_type(obj.default_factory)
+            for k, v in obj.items():
+                result[_asdict_inner(k, dict_factory)] = _asdict_inner(v, dict_factory)
+            return result
+        return obj_type(
+            (_asdict_inner(k, dict_factory), _asdict_inner(v, dict_factory))
+            for k, v in obj.items()
+        )
+    elif issubclass(obj_type, list):
+        return obj_type(_asdict_inner(v, dict_factory) for v in obj)
+    else:
+        return copy.deepcopy(obj)

message_ix_models/util/cache.py

@@ -13,7 +13,7 @@
 import json
 import logging
 from collections.abc import Callable
-from dataclasses import asdict, is_dataclass
+from dataclasses import is_dataclass
 from types import FunctionType
 from typing import TYPE_CHECKING, Union
 
@@ -22,6 +22,7 @@
 import sdmx.model
 import xarray as xr
 
+from ._dataclasses import asdict
 from .context import Context
 from .scenarioinfo import ScenarioInfo
 

message_ix_models/util/config.py

@@ -2,13 +2,14 @@
 import os
 import pickle
 from collections.abc import Mapping, MutableMapping, Sequence
-from dataclasses import asdict, dataclass, field, fields, is_dataclass, replace
+from dataclasses import dataclass, field, fields, is_dataclass, replace
 from hashlib import blake2s
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Hashable, Optional
 
 import ixmp
 
+from ._dataclasses import asdict
 from .scenarioinfo import ScenarioInfo
 
 if TYPE_CHECKING:

message_ix_models/util/context.py

@@ -2,6 +2,7 @@
 
 import logging
 from copy import deepcopy
+from dataclasses import is_dataclass
 from functools import lru_cache
 from importlib import import_module
 from pathlib import Path
@@ -253,7 +254,12 @@ def __setattr__(self, name, value):
     # Particular methods of Context
     def asdict(self) -> dict:
         """Return a :func:`.deepcopy` of the Context's values as a :class:`dict`."""
-        return {k: deepcopy(v) for k, v in self._values.items()}
+        from ._dataclasses import asdict
+
+        return {
+            k: asdict(v) if is_dataclass(v) else deepcopy(v)
+            for k, v in self._values.items()
+        }
 
     def clone_to_dest(self, create=True) -> "message_ix.Scenario":
         """Return a scenario based on the ``--dest`` command-line option.