diff --git a/pkg-py/src/querychat/querychat.py b/pkg-py/src/querychat/querychat.py index 3f6ea616..5b142576 100644 --- a/pkg-py/src/querychat/querychat.py +++ b/pkg-py/src/querychat/querychat.py @@ -224,11 +224,20 @@ def df_to_html(df: IntoFrame, maxrows: int = 5) -> str: HTML string representation of the table """ + # Convert to Narwhals DataFrame if it's not already one if isinstance(df, (nw.LazyFrame, nw.DataFrame)): - df_short = df.lazy().head(maxrows).collect() - nrow_full = df.lazy().select(nw.len()).collect().item() + nw_df = df else: - raise TypeError("df must be a Narwhals DataFrame or LazyFrame") + # Try to convert using nw.from_native (supports pandas and other formats) + try: + nw_df = nw.from_native(df) + except Exception as e: + raise TypeError( + "df must be a Narwhals DataFrame, LazyFrame, or compatible DataFrame (e.g., pandas)", + ) from e + + df_short = nw_df.lazy().head(maxrows).collect() + nrow_full = nw_df.lazy().select(nw.len()).collect().item() # Generate HTML table table_html = df_short.to_pandas().to_html( diff --git a/pkg-py/tests/test_query_function.py b/pkg-py/tests/test_query_function.py new file mode 100644 index 00000000..982efacd --- /dev/null +++ b/pkg-py/tests/test_query_function.py @@ -0,0 +1,116 @@ +import sqlite3 +import tempfile +from pathlib import Path + +import pandas as pd +import pytest +from sqlalchemy import create_engine +from src.querychat.datasource import DataFrameSource, SQLAlchemySource +from src.querychat.querychat import df_to_html + + +@pytest.fixture +def sample_dataframe(): + """Create a sample pandas DataFrame for testing.""" + return pd.DataFrame( + { + "id": [1, 2, 3, 4, 5], + "name": ["Alice", "Bob", "Charlie", "Diana", "Eve"], + "age": [25, 30, 35, 28, 32], + "salary": [50000, 60000, 70000, 55000, 65000], + }, + ) + + +@pytest.fixture +def test_db_engine_with_data(): + """Create a temporary SQLite database with test data.""" + temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db") # noqa: SIM115 + temp_db.close() + + conn = sqlite3.connect(temp_db.name) + cursor = conn.cursor() + + cursor.execute(""" + CREATE TABLE employees ( + id INTEGER PRIMARY KEY, + name TEXT, + age INTEGER, + salary REAL + ) + """) + + test_data = [ + (1, "Alice", 25, 50000), + (2, "Bob", 30, 60000), + (3, "Charlie", 35, 70000), + (4, "Diana", 28, 55000), + (5, "Eve", 32, 65000), + ] + + cursor.executemany( + "INSERT INTO employees (id, name, age, salary) VALUES (?, ?, ?, ?)", + test_data, + ) + + conn.commit() + conn.close() + + engine = create_engine(f"sqlite:///{temp_db.name}") + yield engine + + # Cleanup + Path(temp_db.name).unlink() + + +def test_df_to_html_with_dataframe_source_result(sample_dataframe): + """Test that df_to_html() works with results from DataFrameSource.execute_query().""" + source = DataFrameSource(sample_dataframe, "employees") + + # Execute query to get pandas DataFrame + result_df = source.execute_query("SELECT * FROM employees WHERE age > 25") + + # This should succeed after the fix + html_output = df_to_html(result_df) + + # Verify the HTML contains expected content + assert isinstance(html_output, str) + assert "