fix #108

Author: PaleNeutron

dataframe_image/_pandas_accessor.py

@@ -1,3 +1,4 @@
+import io
 from pathlib import Path
 from typing import Literal
 
@@ -82,7 +83,11 @@ def export(
         from .converter.matplotlib_table import MatplotlibTableConverter
 
         # get extension from filename without dot
-        extension = Path(filename).suffix
+        if isinstance(filename, io.IOBase):
+            extension = "png"
+        else:
+            extension = Path(filename).suffix
+
         if extension.startswith("."):
             extension = extension[1:]
         converter = MatplotlibTableConverter(

tests/test_df_image.py

@@ -1,3 +1,4 @@
+from io import BytesIO
 import random
 import string
 
@@ -34,6 +35,18 @@ def test_styled(document_name, converter, dpi):
         dpi=dpi,
     )
 
+@pytest.mark.parametrize("converter", converters)
+def test_styled2(document_name, converter):
+    col_headers = {
+        "selector": ".col_heading, thead",
+        "props": "color: white; background-color: #1d5632; font-size: 11px"
+        }
+            
+    df = pd.DataFrame(np.random.rand(6, 4))
+    df_styled = df.style.set_table_styles([col_headers])
+    dfi.export(df_styled, f"tests/test_output/{document_name}.png", table_conversion=converter)
+
+
 
 @pytest.mark.parametrize("dpi", test_dpi_values)
 @pytest.mark.parametrize("converter", converters)
@@ -81,3 +94,8 @@ def test_long_column_headers(document_name, converter, dpi):
         dpi=dpi,
         max_rows=-1,
     )
+
+
+def test_save_using_bytesio():
+    buf = BytesIO()
+    df.dfi.export(buf, table_conversion='matplotlib')