From ecf6392fdb7b7f3e73abe8e53f2682c8c35768d1 Mon Sep 17 00:00:00 2001
From: John Lyu <paleneutron@outlook.com>
Date: Sat, 2 Dec 2023 12:15:18 +0800
Subject: [PATCH] fix #108

---
 dataframe_image/_pandas_accessor.py |  7 ++++++-
 tests/test_df_image.py              | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/dataframe_image/_pandas_accessor.py b/dataframe_image/_pandas_accessor.py
index cdccdb1..1725dfa 100644
--- a/dataframe_image/_pandas_accessor.py
+++ b/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(
diff --git a/tests/test_df_image.py b/tests/test_df_image.py
index 80e7dbd..af5babf 100644
--- a/tests/test_df_image.py
+++ b/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')
\ No newline at end of file