output list not array

marl · Jan 8, 2017 · 01b9137 · 01b9137
1 parent 424585d
commit 01b9137
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/medleydb/utils.py b/medleydb/utils.py
@@ -165,8 +165,8 @@ def artist_conditional_split(trackid_list=None, test_size=0.15, num_splits=5,
     splits = []
     for train, test in splitter:
         splits.append({
-            'train': trackid_array[train],
-            'test': trackid_array[test]
+            'train': list(trackid_array[train]),
+            'test': list(trackid_array[test])
         })
 
     return splits

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -119,7 +119,7 @@ def test_trackid_list(self):
         self.assertEqual(1, len(splits))
         self.assertEqual(
             sorted(trackid_list),
-            sorted(list(splits[0]['train']) + list(splits[0]['test']))
+            sorted(splits[0]['train'] + splits[0]['test'])
         )
 
     def test_artist_index(self):
@@ -146,6 +146,6 @@ def test_artist_index(self):
         self.assertEqual(1, len(splits))
         self.assertEqual(
             sorted(trackid_list),
-            sorted(list(splits[0]['train']) + list(splits[0]['test']))
+            sorted(splits[0]['train'] + splits[0]['test'])
         )