Fix to dict conversion of DatasetInfo/Features (#4741)

* Add custom asdict

* Add test

* One more test

* Comment

Author: mariosasko

src/datasets/arrow_dataset.py

@@ -27,7 +27,6 @@
 from collections import Counter, UserDict
 from collections.abc import Mapping
 from copy import deepcopy
-from dataclasses import asdict
 from functools import partial, wraps
 from io import BytesIO
 from math import ceil, floor
@@ -95,7 +94,7 @@
 from .utils._hf_hub_fixes import create_repo
 from .utils.file_utils import _retry, cached_path, estimate_dataset_size, hf_hub_url
 from .utils.info_utils import is_small_dataset
-from .utils.py_utils import convert_file_size_to_int, unique_values
+from .utils.py_utils import asdict, convert_file_size_to_int, unique_values
 from .utils.stratify import stratified_shuffle_split_generate_indices
 from .utils.tf_utils import minimal_tf_collate_fn
 from .utils.typing import PathLike

src/datasets/arrow_writer.py

@@ -16,7 +16,6 @@
 import json
 import os
 import sys
-from dataclasses import asdict
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
@@ -39,7 +38,7 @@
 from .table import array_cast, cast_array_to_feature, table_cast
 from .utils import logging
 from .utils.file_utils import hash_url_to_filename
-from .utils.py_utils import first_non_null_value
+from .utils.py_utils import asdict, first_non_null_value
 
 
 logger = logging.get_logger(__name__)

src/datasets/features/features.py

@@ -19,7 +19,7 @@
 import re
 import sys
 from collections.abc import Iterable, Mapping
-from dataclasses import InitVar, _asdict_inner, dataclass, field, fields
+from dataclasses import InitVar, dataclass, field, fields
 from functools import reduce, wraps
 from operator import mul
 from typing import Any, ClassVar, Dict, List, Optional
@@ -37,7 +37,7 @@
 from .. import config
 from ..table import array_cast
 from ..utils import logging
-from ..utils.py_utils import first_non_null_value, zip_dict
+from ..utils.py_utils import asdict, first_non_null_value, zip_dict
 from .audio import Audio
 from .image import Image, encode_pil_image
 from .translation import Translation, TranslationVariableLanguages
@@ -1598,7 +1598,7 @@ def from_dict(cls, dic) -> "Features":
         return cls(**obj)
 
     def to_dict(self):
-        return _asdict_inner(self, dict)
+        return asdict(self)
 
     def encode_example(self, example):
         """

src/datasets/fingerprint.py

@@ -5,7 +5,6 @@
 import shutil
 import tempfile
 import weakref
-from dataclasses import asdict
 from functools import wraps
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
@@ -19,7 +18,7 @@
 from .table import ConcatenationTable, InMemoryTable, MemoryMappedTable, Table
 from .utils.deprecation_utils import deprecated
 from .utils.logging import get_logger
-from .utils.py_utils import dumps
+from .utils.py_utils import asdict, dumps
 
 
 if TYPE_CHECKING:

src/datasets/info.py

@@ -32,7 +32,7 @@
 import dataclasses
 import json
 import os
-from dataclasses import asdict, dataclass, field
+from dataclasses import dataclass, field
 from typing import Dict, List, Optional, Union
 
 from . import config
@@ -41,7 +41,7 @@
 from .tasks import TaskTemplate, task_template_from_dict
 from .utils import Version
 from .utils.logging import get_logger
-from .utils.py_utils import unique_values
+from .utils.py_utils import asdict, unique_values
 
 
 logger = get_logger(__name__)

src/datasets/utils/py_utils.py

@@ -18,6 +18,7 @@
 """
 
 import contextlib
+import copy
 import functools
 import itertools
 import os
@@ -26,6 +27,7 @@
 import sys
 import types
 from contextlib import contextmanager
+from dataclasses import fields, is_dataclass
 from io import BytesIO as StringIO
 from multiprocessing import Pool, RLock
 from shutil import disk_usage
@@ -151,6 +153,41 @@ def string_to_dict(string: str, pattern: str) -> Dict[str, str]:
     return _dict
 
 
+def asdict(obj):
+    """Convert an object to its dictionary representation recursively."""
+
+    # Implementation based on https://docs.python.org/3/library/dataclasses.html#dataclasses.asdict
+
+    def _is_dataclass_instance(obj):
+        # https://docs.python.org/3/library/dataclasses.html#dataclasses.is_dataclass
+        return is_dataclass(obj) and not isinstance(obj, type)
+
+    def _asdict_inner(obj):
+        if _is_dataclass_instance(obj):
+            result = {}
+            for f in fields(obj):
+                value = _asdict_inner(getattr(obj, f.name))
+                result[f.name] = value
+            return result
+        elif isinstance(obj, tuple) and hasattr(obj, "_fields"):
+            # obj is a namedtuple
+            return type(obj)(*[_asdict_inner(v) for v in obj])
+        elif isinstance(obj, (list, tuple)):
+            # Assume we can create an object of this type by passing in a
+            # generator (which is not true for namedtuples, handled
+            # above).
+            return type(obj)(_asdict_inner(v) for v in obj)
+        elif isinstance(obj, dict):
+            return {_asdict_inner(k): _asdict_inner(v) for k, v in obj.items()}
+        else:
+            return copy.deepcopy(obj)
+
+    if not isinstance(obj, dict) and not _is_dataclass_instance(obj):
+        raise TypeError(f"{obj} is not a dict or a dataclass")
+
+    return _asdict_inner(obj)
+
+
 @contextlib.contextmanager
 def temporary_assignment(obj, attr, value):
     """Temporarily assign obj.attr to value."""

tests/features/test_features.py

@@ -1,5 +1,4 @@
 import datetime
-from dataclasses import asdict
 from unittest import TestCase
 from unittest.mock import patch
 
@@ -20,6 +19,7 @@
 )
 from datasets.features.translation import Translation, TranslationVariableLanguages
 from datasets.info import DatasetInfo
+from datasets.utils.py_utils import asdict
 
 from ..utils import require_jax, require_tf, require_torch
 
@@ -101,6 +101,13 @@ def test_feature_named_type(self):
         reloaded_features = Features.from_dict(asdict(ds_info)["features"])
         assert features == reloaded_features
 
+    def test_class_label_feature_with_no_labels(self):
+        """reference: issue #4681"""
+        features = Features({"label": ClassLabel(names=[])})
+        ds_info = DatasetInfo(features=features)
+        reloaded_features = Features.from_dict(asdict(ds_info)["features"])
+        assert features == reloaded_features
+
     def test_reorder_fields_as(self):
         features = Features(
             {

tests/test_py_utils.py

@@ -1,9 +1,10 @@
+from dataclasses import dataclass
 from unittest import TestCase
 
 import numpy as np
 import pytest
 
-from datasets.utils.py_utils import NestedDataStructure, map_nested, temp_seed, temporary_assignment, zip_dict
+from datasets.utils.py_utils import NestedDataStructure, asdict, map_nested, temp_seed, temporary_assignment, zip_dict
 
 from .utils import require_tf, require_torch
 
@@ -16,6 +17,12 @@ def add_one(i):  # picklable for multiprocessing
     return i + 1
 
 
+@dataclass
+class A:
+    x: int
+    y: str
+
+
 class PyUtilsTest(TestCase):
     def test_map_nested(self):
         s1 = {}
@@ -175,3 +182,16 @@ def test_nested_data_structure_data(input_data):
 def test_flatten(data, expected_output):
     output = NestedDataStructure(data).flatten()
     assert output == expected_output
+
+
+def test_asdict():
+    input = A(x=1, y="foobar")
+    expected_output = {"x": 1, "y": "foobar"}
+    assert asdict(input) == expected_output
+
+    input = {"a": {"b": A(x=10, y="foo")}, "c": [A(x=20, y="bar")]}
+    expected_output = {"a": {"b": {"x": 10, "y": "foo"}}, "c": [{"x": 20, "y": "bar"}]}
+    assert asdict(input) == expected_output
+
+    with pytest.raises(TypeError):
+        asdict([1, A(x=10, y="foo")])