Skip to content

Commit

Permalink
Backport dataclasses.asdict(…) for Python ≤ 3.11
Browse files Browse the repository at this point in the history
Work around python/cpython#79721 on these versions.
  • Loading branch information
khaeru committed Dec 3, 2024
1 parent b8d9622 commit efc5ef8
Show file tree
Hide file tree
Showing 5 changed files with 136 additions and 5 deletions.
11 changes: 8 additions & 3 deletions message_ix_models/tests/util/test_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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."""
Expand Down
28 changes: 28 additions & 0 deletions message_ix_models/tests/util/test_scenarioinfo.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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()
Expand Down
95 changes: 95 additions & 0 deletions message_ix_models/util/_dataclasses.py
Original file line number Diff line number Diff line change
@@ -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)
3 changes: 2 additions & 1 deletion message_ix_models/util/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
4 changes: 3 additions & 1 deletion message_ix_models/util/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,9 @@ 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) 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.
Expand Down

0 comments on commit efc5ef8

Please sign in to comment.