querychat 0.2.0__tar.gz → 0.2.2__tar.gz

{querychat-0.2.0 → querychat-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: querychat
-Version: 0.2.0
+Version: 0.2.2
 Summary: Chat with your data using natural language
 Project-URL: Homepage, https://github.com/posit-dev/querychat
 Project-URL: Repository, https://github.com/posit-dev/querychat

{querychat-0.2.0 → querychat-0.2.2}/pkg-py/src/querychat/datasource.py

@@ -3,12 +3,13 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, ClassVar, Protocol
 
 import duckdb
-import narwhals as nw
+import narwhals.stable.v1 as nw
 import pandas as pd
 from sqlalchemy import inspect, text
 from sqlalchemy.sql import sqltypes
 
 if TYPE_CHECKING:
+    from narwhals.stable.v1.typing import IntoFrame
     from sqlalchemy.engine import Connection, Engine
 
 
@@ -58,8 +59,9 @@ class DataFrameSource:
     """A DataSource implementation that wraps a pandas DataFrame using DuckDB."""
 
     db_engine: ClassVar[str] = "DuckDB"
+    _df: nw.DataFrame | nw.LazyFrame
 
-    def __init__(self, df: pd.DataFrame, table_name: str):
+    def __init__(self, df: IntoFrame, table_name: str):
         """
         Initialize with a pandas DataFrame.
 
@@ -69,9 +71,10 @@ class DataFrameSource:
 
         """
         self._conn = duckdb.connect(database=":memory:")
-        self._df = df
+        self._df = nw.from_native(df)
         self._table_name = table_name
-        self._conn.register(table_name, df)
+        # TODO(@gadenbuie): If the data frame is already SQL-backed, maybe we shouldn't be making a new copy here.
+        self._conn.register(table_name, self._df.lazy().collect().to_pandas())
 
     def get_schema(self, *, categorical_threshold: int) -> str:
         """
@@ -86,10 +89,15 @@ class DataFrameSource:
             String describing the schema
 
         """
-        ndf = nw.from_native(self._df)
-
         schema = [f"Table: {self._table_name}", "Columns:"]
 
+        # Ensure we're working with a DataFrame, not a LazyFrame
+        ndf = (
+            self._df.head(10).collect()
+            if isinstance(self._df, nw.LazyFrame)
+            else self._df
+        )
+
         for column in ndf.columns:
             # Map pandas dtypes to SQL-like types
             dtype = ndf[column].dtype
@@ -149,7 +157,8 @@ class DataFrameSource:
             The complete dataset as a pandas DataFrame
 
         """
-        return self._df.copy()
+        # TODO(@gadenbuie): This should just return `self._df` and not a pandas DataFrame
+        return self._df.lazy().collect().to_pandas()
 
 
 class SQLAlchemySource:

{querychat-0.2.0 → querychat-0.2.2}/pkg-py/src/querychat/querychat.py

@@ -18,7 +18,7 @@ from typing import (
 
 import chatlas
 import chevron
-import narwhals as nw
+import narwhals.stable.v1 as nw
 import sqlalchemy
 from shiny import Inputs, Outputs, Session, module, reactive, ui
 
@@ -26,7 +26,7 @@ from ._utils import temp_env_vars
 
 if TYPE_CHECKING:
     import pandas as pd
-    from narwhals.typing import IntoFrame
+    from narwhals.stable.v1.typing import IntoFrame
 
 from .datasource import DataFrameSource, DataSource, SQLAlchemySource
 
@@ -224,16 +224,15 @@ def df_to_html(df: IntoFrame, maxrows: int = 5) -> str:
         HTML string representation of the table
 
     """
-    df_short: nw.DataFrame[Any]
-
-    if isinstance(df, nw.LazyFrame):
-        ndf_eager = df.collect()
-        df_short = df.head(maxrows).collect()
-    elif isinstance(df, nw.DataFrame):
-        ndf_eager = df
-        df_short = df.head(maxrows)
+    ndf = nw.from_native(df)
+
+    if isinstance(ndf, (nw.LazyFrame, nw.DataFrame)):
+        df_short = ndf.lazy().head(maxrows).collect()
+        nrow_full = ndf.lazy().select(nw.len()).collect().item()
     else:
-        raise TypeError("df must be a Narwhals DataFrame or LazyFrame")
+        raise TypeError(
+            "Must be able to convert `df` into a Narwhals DataFrame or LazyFrame",
+        )
 
     # Generate HTML table
     table_html = df_short.to_pandas().to_html(
@@ -242,9 +241,9 @@ def df_to_html(df: IntoFrame, maxrows: int = 5) -> str:
     )
 
     # Add note about truncated rows if needed
-    if len(df_short) != len(ndf_eager):
+    if len(df_short) != nrow_full:
         rows_notice = (
-            f"\n\n(Showing only the first {maxrows} rows out of {len(ndf_eager)}.)\n"
+            f"\n\n(Showing only the first {maxrows} rows out of {nrow_full}.)\n"
         )
     else:
         rows_notice = ""
@@ -415,15 +414,12 @@ def init(
     data_source_obj: DataSource
     if isinstance(data_source, sqlalchemy.Engine):
         data_source_obj = SQLAlchemySource(data_source, table_name)
-    elif isinstance(data_source, (nw.DataFrame, nw.LazyFrame)):
+    else:
         data_source_obj = DataFrameSource(
-            nw.to_native(data_source),
+            data_source,
             table_name,
         )
-    else:
-        raise TypeError(
-            "`data_source` must be a Narwhals DataFrame or LazyFrame, or a SQLAlchemy Engine",
-        )
+
     # Process greeting
     if greeting is None:
         print(

{querychat-0.2.0 → querychat-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "querychat"
-version = "0.2.0"
+version = "0.2.2"
 description = "Chat with your data using natural language"
 readme = "pkg-py/README.md"
 requires-python = ">=3.9"
@@ -124,6 +124,7 @@ extend-ignore = [
     "D107", # Missing docstring in __init__
     "D205", # 1 blank line required between summary line and description
     "UP045", # Use `X | NULL` for type annotations, not `Optional[X]`
+    "TD003", # TODO doesn't need to have an issue link
 ]
 extend-select = [
     # "C90", # C90; mccabe: https://docs.astral.sh/ruff/rules/complex-structure/