fix test

Author: Martin Durant

streamz/dataframe/tests/test_dataframe_utils.py

@@ -18,14 +18,14 @@ def test_utils_get_base_frame_type_pandas():
     with pytest.raises(TypeError):
         get_base_frame_type("Index", is_index_like, df)
 
-    with pytest.raises(TypeError):
-        get_base_frame_type("DataFrame", is_dataframe_like, df.x)
+    # casts Series to DataFrame, if that's what we ask for
+    assert pd.DataFrame == get_base_frame_type("DataFrame", is_dataframe_like, df.x)
     assert pd.Series == get_base_frame_type("Series", is_series_like, df.x)
     with pytest.raises(TypeError):
         get_base_frame_type("Index", is_index_like, df.x)
 
-    with pytest.raises(TypeError):
-        get_base_frame_type("DataFrame", is_dataframe_like, df.index)
+    # casts Series to DataFrame, if that's what we ask for
+    assert pd.DataFrame == get_base_frame_type("DataFrame", is_dataframe_like, df.index)
     with pytest.raises(TypeError):
         get_base_frame_type("Series", is_series_like, df.index)
     assert issubclass(get_base_frame_type("Index", is_index_like, df.index), pd.Index)

streamz/dataframe/utils.py

@@ -38,14 +38,14 @@ def get_base_frame_type(frame_name, is_frame_like, example=None):
        Returns the base type of streaming objects if type checks pass."""
     if example is None:
         raise TypeError("Missing required argument:'example'")
-    if not is_frame_like(example):
-        try:
-            import pandas as pd
-            example = pd.DataFrame(example)
-        except (TypeError, ValueError) as e:
-            msg = "Streaming {0} expects an example of {0} like objects. Got: {1}."\
-                                                 .format(frame_name, example)
-            raise TypeError(msg) from e
+    if is_frame_like is is_dataframe_like and not is_frame_like(example):
+        import pandas as pd
+        example = pd.DataFrame(example)
+
+    elif not is_frame_like(example):
+        msg = "Streaming {0} expects an example of {0} like objects. Got: {1}."\
+                                             .format(frame_name, example)
+        raise TypeError(msg)
     return type(example)