fix: avoid casting tuples after Dataset.map (#4993)

* fix: avoid casting tuples after Dataset.map * fix: fix test_cast_to_python_objects_tuple test
huggingface · Sep 20, 2022 · 8ba0522 · 8ba0522 · github-actions · Sep 20, 2022
1 parent 1b4c3cb
commit 8ba0522
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/datasets/features/features.py b/src/datasets/features/features.py
@@ -379,12 +379,12 @@ def _cast_to_python_objects(obj: Any, only_1d_for_numpy: bool, optimize_list_cas
                     for elmt in obj
                 ], True
             else:
-                if isinstance(obj, list):
+                if isinstance(obj, (list, tuple)):
                     return obj, False
                 else:
                     return list(obj), True
         else:
-            return obj if isinstance(obj, list) else [], isinstance(obj, tuple)
+            return obj, False
     else:
         return obj, False
 

diff --git a/tests/features/test_features.py b/tests/features/test_features.py
@@ -463,7 +463,7 @@ def test_cast_to_python_objects_list(self):
 
     def test_cast_to_python_objects_tuple(self):
         obj = {"col_1": [{"vec": (1, 2, 3), "txt": "foo"}] * 3, "col_2": [(1, 2), (3, 4), (5, 6)]}
-        expected_obj = {"col_1": [{"vec": [1, 2, 3], "txt": "foo"}] * 3, "col_2": [[1, 2], [3, 4], [5, 6]]}
+        expected_obj = {"col_1": [{"vec": (1, 2, 3), "txt": "foo"}] * 3, "col_2": [(1, 2), (3, 4), (5, 6)]}
         casted_obj = cast_to_python_objects(obj)
         self.assertDictEqual(casted_obj, expected_obj)