fix: filter sql defs from generated return (#6805)

## 📝 Summary

Previously defining a sql variable would also add it to the return
signature. This is incorrect, since one would actually get a name error.

Added a test to prevent future regressions.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: dmadisetti

examples/sql/histograms.py

@@ -12,7 +12,7 @@
 
 import marimo
 
-__generated_with = "0.15.5"
+__generated_with = "0.17.0"
 app = marimo.App(width="medium")
 
 
@@ -60,7 +60,7 @@ def _(URL, mo):
         FROM dataset
         """
     )
-    return (dataset,)
+    return
 
 
 @app.cell

examples/sql/misc/electric_vehicles.py

@@ -11,7 +11,7 @@
 
 import marimo
 
-__generated_with = "0.16.0"
+__generated_with = "0.17.0"
 app = marimo.App(width="medium")
 
 
@@ -36,7 +36,7 @@ def _(mo):
         select * from evs
         """
     )
-    return (evs,)
+    return
 
 
 @app.cell

examples/sql/misc/sql_cars.py

@@ -1,18 +1,18 @@
 # /// script
 # requires-python = ">=3.9"
 # dependencies = [
-#     "altair==5.4.1",
-#     "duckdb==1.1.1",
+#     "altair>=5.4.1",
+#     "duckdb>=1.1.1",
 #     "marimo",
-#     "polars==1.18.0",
-#     "pyarrow==18.1.0",
-#     "vega-datasets==0.9.0",
+#     "polars>=1.18.0",
+#     "pyarrow>=18.1.0",
+#     "vega-datasets>=0.9.0",
 # ]
 # ///
 
 import marimo
 
-__generated_with = "0.16.0"
+__generated_with = "0.17.0"
 app = marimo.App(width="medium")
 
 
@@ -31,7 +31,7 @@ def _(cars_df, mo):
         CREATE OR REPLACE TABLE cars AS SELECT * FROM cars_df;
         """
     )
-    return (cars,)
+    return
 
 
 @app.cell
@@ -54,7 +54,7 @@ def _(origin):
 @app.cell
 def _(mo, origin, top_n):
     mo.md(
-        f"""##Top {top_n.value} Cars {f"in {origin.value}" if origin.value != None else ""} """
+        f"""##Top {top_n.value} Cars {f"in {origin.value}" if origin.value != None else ""}"""
     )
     return
 

examples/sql/read_csv.py

@@ -11,7 +11,7 @@
 
 import marimo
 
-__generated_with = "0.16.0"
+__generated_with = "0.17.0"
 app = marimo.App(width="medium")
 
 
@@ -118,7 +118,7 @@ def _(mo):
         CREATE TABLE myTable AS SELECT * FROM "data.csv"
         """
     )
-    return (mytable,)
+    return
 
 
 @app.cell

examples/sql/read_json.py

@@ -11,7 +11,7 @@
 
 import marimo
 
-__generated_with = "0.16.0"
+__generated_with = "0.17.0"
 app = marimo.App(width="medium")
 
 
@@ -128,7 +128,7 @@ def _(mo):
         CREATE OR REPLACE TABLE myTable AS SELECT * FROM 'data.json'
         """
     )
-    return (mytable,)
+    return
 
 
 @app.cell

examples/sql/read_parquet.py

@@ -11,7 +11,7 @@
 
 import marimo
 
-__generated_with = "0.16.0"
+__generated_with = "0.17.0"
 app = marimo.App(width="medium")
 
 
@@ -117,7 +117,7 @@ def _(mo):
         CREATE OR REPLACE TABLE myTable AS SELECT * FROM 'data.parquet'
         """
     )
-    return (mytable,)
+    return
 
 
 @app.cell

examples/ui/table.py

@@ -1,6 +1,6 @@
 import marimo
 
-__generated_with = "0.16.0"
+__generated_with = "0.17.0"
 app = marimo.App()
 
 
@@ -43,14 +43,15 @@ def cell_hover(row_id: str, column_name: str, value) -> str:
         hover_template=cell_hover,
     )
     hover_table
-    return (hover_table,)
+    return
 
 
 @app.cell
 def _(table):
     table.value
     return
 
+
 @app.cell
 def _(mo):
     # Demonstrate a long table with a sticky header and a custom max height
@@ -61,8 +62,7 @@ def _(mo):
         max_height=300,
     )
     long_table
-    return (long_table,)
-
+    return
 
 
 if __name__ == "__main__":

marimo/_ast/codegen.py

@@ -236,15 +236,24 @@ def to_functiondef(
     # other static analysis tools can capture unused variables across cells.
     defs: tuple[str, ...] = tuple()
     if cell.defs:
+        # SQL defs should not be included in the return value.
+        sql_defs = (
+            {
+                name
+                for name, value in variable_data.items()
+                if value.language == "sql"
+            }
+            if variable_data
+            else set()
+        )
         # There are possible name error cases where a cell defines, and also
         # requires a variable. We remove defs from the signature such that
         # this causes a lint error in pyright.
-        if used_refs is None:
-            defs = tuple(name for name in sorted(cell.defs))
-        else:
-            defs = tuple(
-                name for name in sorted(cell.defs) if name in used_refs
-            )
+        defs = tuple(
+            name for name in sorted(cell.defs) if name not in sql_defs
+        )
+        if used_refs is not None:
+            defs = tuple(name for name in defs if name in used_refs)
 
     decorator = to_decorator(cell.config, fn=fn)
     prefix = "" if not cell.is_coroutine() else "async "

tests/_ast/test_codegen.py

@@ -662,6 +662,34 @@ def test_dotted_names_filtered_from_signature(self) -> None:
         assert "my_schema.pokemon_db" not in fndef
         assert "def foo(my_schema.pokemon_db" not in fndef
 
+    def test_sql_defs_filtered_from_return(self) -> None:
+        """Test that SQL definitions are filtered from return but can still be referenced."""
+
+        # Cell 1: defines a SQL variable (cars) - should NOT be in return
+        code1 = "empty = mo.sql('CREATE TABLE cars_df ();')"
+        # Cell 2: uses the SQL variable (cars) - should appear in signature
+        code2 = "result = cars_df.filter(lambda x: x > 0); empty"
+        expected = wrap_generate_filecontents(
+            [code1, code2], ["cell1", "cell2"]
+        )
+        assert (
+            "\n".join(
+                [
+                    "@app.cell",
+                    "def cell1(mo):",
+                    "    empty = mo.sql('CREATE TABLE cars_df ();')",
+                    "    return (empty,)",  # Doesn't return cars_df
+                    "",
+                    "",
+                    "@app.cell",
+                    "def cell2(cars_df, empty):",
+                    "    result = cars_df.filter(lambda x: x > 0); empty",
+                    "    return",
+                ]
+            )
+            in expected
+        )
+
     def test_should_remove_defaults(self) -> None:
         code = "x = 0"
         cell = compile_cell(code)