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_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_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)), } 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_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": [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_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_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": [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_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_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