def test_cast_to_python_objects_series(self):
obj = {
"col_1": pd.Series([{"vec": [1, 2, 3], "txt": "foo"}] * 3),
"col_2": pd.Series([[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)
def test_cast_to_python_or_numpy(self):
obj = {"col_1": [{"vec": np.arange(1, 4), "txt": "foo"}] * 3, "col_2": np.arange(1, 7).reshape(3, 2)}
expected_obj = {
"col_1": [{"vec": np.array([1, 2, 3]), "txt": "foo"}] * 3,
"col_2": np.array([[1, 2], [3, 4], [5, 6]]),
}
casted_obj = cast_to_python_objects(obj)
dict_diff(casted_obj, expected_obj)
def test_cast_to_python_objects_tf(self):
import tensorflow as tf
obj = {
"col_1": [{"vec": tf.constant(np.arange(1, 4)), "txt": "foo"}] * 3,
"col_2": tf.constant(np.arange(1, 7).reshape(3, 2)),
}
expected_obj = {
"col_1": [{"vec": np.array([1, 2, 3]), "txt": "foo"}] * 3,
"col_2": np.array([[1, 2], [3, 4], [5, 6]]),
}
casted_obj = cast_to_python_objects(obj)
dict_diff(casted_obj, expected_obj)
def test_cast_to_python_objects_jax(self):
import jax.numpy as jnp
obj = {
"col_1": [{"vec": jnp.array(np.arange(1, 4)), "txt": "foo"}] * 3,
"col_2": jnp.array(np.arange(1, 7).reshape(3, 2)),
}
assert obj["col_2"].dtype == jnp.int32
expected_obj = {
"col_1": [{"vec": np.array([1, 2, 3], dtype=np.int32), "txt": "foo"}] * 3,
"col_2": np.array([[1, 2], [3, 4], [5, 6]], dtype=np.int32),
}
casted_obj = cast_to_python_objects(obj)
dict_diff(casted_obj, expected_obj)
def test_cast_to_python_objects_dataframe(self):
obj = pd.DataFrame({
"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)
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]]
}
casted_obj = cast_to_python_objects(obj)
self.assertDictEqual(casted_obj, expected_obj)
def test_dont_iterate_over_each_element_in_a_list(self, mocked_cast):
obj = {"col_1": [[1, 2], [3, 4], [5, 6]]}
cast_to_python_objects(obj)
self.assertEqual(mocked_cast.call_count, 4) # 4 = depth of obj