feat: add visualization layer with altair

app.py

@@ -28,6 +28,7 @@
 
 # Import authentication
 from src.simple_auth_components import simple_auth_wrapper
+from src.visualization import render_visualization
 
 # Configure page with professional styling
 favicon_path = get_favicon_path()
@@ -361,6 +362,8 @@ def display_results(result_df: pd.DataFrame, title: str, execution_time: float =
         height = min(600, max(200, len(result_df) * 35 + 50))  # Dynamic height based on rows
         st.dataframe(result_df, use_container_width=True, height=height)
 
+        render_visualization(result_df)
+
         st.markdown("</div>", unsafe_allow_html=True)
 
     else:

requirements.txt

@@ -1,5 +1,6 @@
 # Core application dependencies
 streamlit>=1.40.0
+altair>=5.3.0
 duckdb>=0.9.0
 pandas>=2.2.0
 numpy>=1.26.4,<2.0

src/visualization.py

@@ -0,0 +1,97 @@
+
+import streamlit as st
+import pandas as pd
+import altair as alt
+
+def make_chart(df: pd.DataFrame, chart_type: str, x: str, y: str, color: str | None = None):
+    """
+    Create an Altair chart based on the given parameters.
+    """
+    if chart_type == "Bar":
+        chart = alt.Chart(df).mark_bar().encode(
+            x=x,
+            y=y,
+        )
+    elif chart_type == "Line":
+        chart = alt.Chart(df).mark_line().encode(
+            x=x,
+            y=y,
+        )
+    elif chart_type == "Scatter":
+        chart = alt.Chart(df).mark_circle().encode(
+            x=x,
+            y=y,
+        )
+    elif chart_type == "Histogram":
+        chart = alt.Chart(df).mark_bar().encode(
+            x=alt.X(x, bin=True),
+            y='count()',
+        )
+    elif chart_type == "Heatmap":
+        chart = alt.Chart(df).mark_rect().encode(
+            x=x,
+            y=y,
+        )
+    else:
+        st.error("Invalid chart type")
+        return None
+
+    if color:
+        chart = chart.encode(color=color)
+
+    return chart.properties(width="container")
+
+def get_chart_recommendation(df: pd.DataFrame) -> tuple[str | None, str | None, str | None]:
+    """
+    Recommend a chart type and axes based on the DataFrame schema.
+    """
+    cols = df.columns
+    numeric_cols = df.select_dtypes(include=['number']).columns
+    categorical_cols = df.select_dtypes(include=['object']).columns
+    datetime_cols = df.select_dtypes(include=['datetime']).columns
+
+    if len(categorical_cols) == 1 and len(numeric_cols) > 1:
+        return "Bar", categorical_cols[0], numeric_cols[0]
+    elif len(numeric_cols) == 1 and len(categorical_cols) == 1:
+        return "Bar", categorical_cols[0], numeric_cols[0]
+    elif len(numeric_cols) == 1 and len(datetime_cols) == 1:
+        return "Line", datetime_cols[0], numeric_cols[0]
+    elif len(numeric_cols) == 2:
+        return "Scatter", numeric_cols[0], numeric_cols[1]
+    elif len(numeric_cols) > 2 and len(categorical_cols) == 0:
+        return "Heatmap", numeric_cols[0], numeric_cols[1]
+
+    return None, None, None
+
+def render_visualization(df: pd.DataFrame):
+    """
+    Render the visualization layer.
+    """
+    st.write("### Visualization")
+
+    chart_type, x_axis, y_axis = get_chart_recommendation(df)
+
+    if chart_type:
+        st.write(f"Recommended Chart: **{chart_type}**")
+
+    cols = df.columns
+    chart_type_options = ["Bar", "Line", "Scatter", "Histogram", "Heatmap"]
+
+    selected_chart_type = st.selectbox("Chart type", chart_type_options, index=chart_type_options.index(chart_type) if chart_type else 0)
+
+    x_axis_options = cols
+    selected_x_axis = st.selectbox("X-axis", x_axis_options, index=x_axis_options.get_loc(x_axis) if x_axis else 0)
+
+    y_axis_options = cols
+    selected_y_axis = st.selectbox("Y-axis", y_axis_options, index=y_axis_options.get_loc(y_axis) if y_axis else 1)
+
+    color_options = [None] + list(cols)
+    selected_color = st.selectbox("Color / Group by", color_options, index=0)
+
+    try:
+        chart = make_chart(df, selected_chart_type, selected_x_axis, selected_y_axis, selected_color)
+        if chart:
+            st.altair_chart(chart, use_container_width=True)
+    except Exception as e:
+        st.warning("Failed to generate chart. Please select compatible columns.")
+        st.dataframe(df)

tests/unit/test_visualization.py

@@ -0,0 +1,80 @@
+
+import pandas as pd
+import pytest
+from src.visualization import make_chart, get_chart_recommendation
+
+@pytest.fixture
+def sample_df():
+    return pd.DataFrame({
+        'A': [1, 2, 3],
+        'B': [4, 5, 6],
+        'C': ['X', 'Y', 'Z'],
+        'D': pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03'])
+    })
+
+def test_make_chart(sample_df):
+    chart = make_chart(sample_df, 'Bar', 'C', 'A')
+    assert chart is not None
+    assert chart.mark == 'bar'
+
+    chart = make_chart(sample_df, 'Line', 'D', 'A')
+    assert chart is not None
+    assert chart.mark == 'line'
+
+    chart = make_chart(sample_df, 'Scatter', 'A', 'B')
+    assert chart is not None
+    assert chart.mark == 'circle'
+
+    chart = make_chart(sample_df, 'Histogram', 'A', 'count()')
+    assert chart is not None
+    assert chart.mark == 'bar'
+
+    chart = make_chart(sample_df, 'Heatmap', 'A', 'B')
+    assert chart is not None
+    assert chart.mark == 'rect'
+
+    chart = make_chart(sample_df, 'Invalid', 'A', 'B')
+    assert chart is None
+
+def test_get_chart_recommendation():
+    # Test case 1: 1 numeric, 1 categorical
+    df1 = pd.DataFrame({'A': [1, 2, 3], 'B': ['X', 'Y', 'Z']})
+    chart_type, x, y = get_chart_recommendation(df1)
+    assert chart_type == 'Bar'
+    assert x == 'B'
+    assert y == 'A'
+
+    # Test case 2: 1 numeric, 1 datetime
+    df2 = pd.DataFrame({'A': [1, 2, 3], 'B': pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03'])})
+    chart_type, x, y = get_chart_recommendation(df2)
+    assert chart_type == 'Line'
+    assert x == 'B'
+    assert y == 'A'
+
+    # Test case 3: 2 numeric
+    df3 = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6]})
+    chart_type, x, y = get_chart_recommendation(df3)
+    assert chart_type == 'Scatter'
+    assert x == 'A'
+    assert y == 'B'
+
+    # Test case 4: >2 numeric, 0 categorical
+    df4 = pd.DataFrame({'A': [1, 2, 3], 'B': [4, 5, 6], 'C': [7, 8, 9]})
+    chart_type, x, y = get_chart_recommendation(df4)
+    assert chart_type == 'Heatmap'
+    assert x == 'A'
+    assert y == 'B'
+
+    # Test case 5: 1 categorical, >1 numeric
+    df5 = pd.DataFrame({'A': ['X', 'Y', 'Z'], 'B': [1, 2, 3], 'C': [4, 5, 6]})
+    chart_type, x, y = get_chart_recommendation(df5)
+    assert chart_type == 'Bar'
+    assert x == 'A'
+    assert y == 'B'
+
+    # Test case 6: No recommendation
+    df6 = pd.DataFrame({'A': ['X', 'Y', 'Z'], 'B': ['a', 'b', 'c']})
+    chart_type, x, y = get_chart_recommendation(df6)
+    assert chart_type is None
+    assert x is None
+    assert y is None