Merge pull request engarde-dev#14 from TomAugspurger/ignore-dim

Ignore dim
topper-123 · Jul 10, 2015 · 5c10817 · 5c10817
2 parents 63ec649 + 529ea9b
commit 5c10817
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/engarde/checks.py b/engarde/checks.py
@@ -98,7 +98,8 @@ def is_shape(df, shape):
     df : DataFrame
     """
     try:
-        assert df.shape == shape
+        check = np.all(np.equal(df.shape, shape) | np.equal(shape, [-1, -1]))
+        assert check
     except AssertionError as e:
         msg = ("Expected shape: {}\n"
                "\t\tActual shape:   {}".format(shape, df.shape))

diff --git a/tests/test_checks.py b/tests/test_checks.py
@@ -138,10 +138,15 @@ def test_monotonic_items():
 
 def test_is_shape():
     shape = 10, 2
+    ig_0 = -1, 2
+    ig_1 = 10, -1
+    shapes = [shape, ig_0, ig_1]
     df = pd.DataFrame(np.random.randn(*shape))
-    tm.assert_frame_equal(df, ck.is_shape(df, shape))
-    result = dc.is_shape(shape=shape)(_add_n)(df)
-    tm.assert_frame_equal(result, df + 1)
+    for shp in shapes:
+        tm.assert_frame_equal(df, ck.is_shape(df, shp))
+    for shp in shapes:
+        result = dc.is_shape(shape=shp)(_add_n)(df)
+        tm.assert_frame_equal(result, df + 1)
 
     with pytest.raises(AssertionError):
         ck.is_shape(df, (9, 2))