remove prints and add test

arviz-devs · Oct 12, 2022 · 3bf9e70 · 3bf9e70
1 parent 8cd1d03
commit 3bf9e70
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 11 deletions.
diff --git a/arviz/data/base.py b/arviz/data/base.py
@@ -228,10 +228,6 @@ def numpy_to_data_array(
     else:
         ary = utils.one_de(ary)
 
-    print(default_dims)
-    print(dims)
-    print(coords)
-    print(ary.shape)
     dims, coords = generate_dims_coords(
         ary.shape[len(default_dims) :],
         var_name,
@@ -241,9 +237,6 @@ def numpy_to_data_array(
         index_origin=index_origin,
         skip_event_dims=skip_event_dims,
     )
-    print("---")
-    print(dims)
-    print(coords)
 
     # reversed order for default dims: 'chain', 'draw'
     if "draw" not in dims and "draw" in default_dims:

diff --git a/arviz/tests/base_tests/test_data.py b/arviz/tests/base_tests/test_data.py
@@ -155,20 +155,21 @@ def test_dims_coords():
     assert len(coords["x_dim_2"]) == 5
 
 
-def test_dims_coords_default_dims():
+@pytest.mark.parametrize("in_dims", (["dim1", "dim2"], ["draw", "dim1", "dim2"], ["chain", "draw", "dim1", "dim2"]))
+def test_dims_coords_default_dims(in_dims):
     shape = 4, 7
     var_name = "x"
     dims, coords = generate_dims_coords(
         shape,
         var_name,
-        dims=["dim1", "dim2"],
+        dims=in_dims,
         coords={"chain": ["a", "b", "c"]},
         default_dims=["chain", "draw"],
     )
     assert "dim1" in dims
     assert "dim2" in dims
-    assert "chain" not in dims
-    assert "draw" not in dims
+    assert ("chain" in dims) == ("chain" in in_dims)
+    assert ("draw" in dims) == ("draw" in in_dims)
     assert len(coords["dim1"]) == 4
     assert len(coords["dim2"]) == 7
     assert len(coords["chain"]) == 3