From dd03b875b573f343b6434b2280a8c8aa80cafd3f Mon Sep 17 00:00:00 2001
From: Dylan Madisetti <dylan@madisetti.me>
Date: Tue, 18 Feb 2025 13:29:38 -0500
Subject: [PATCH] fix: semicolon supression (#3830)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## 📝 Summary

Replaces #3806

## 🔍 Description of Changes

Resolves #3726

@mscolnick

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 marimo/_ast/compiler.py     | 26 +++++++++++-
 tests/_ast/test_compiler.py | 80 +++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+), 1 deletion(-)

diff --git a/marimo/_ast/compiler.py b/marimo/_ast/compiler.py
index b720a14ea61..b92e284e76b 100644
--- a/marimo/_ast/compiler.py
+++ b/marimo/_ast/compiler.py
@@ -48,6 +48,28 @@ def get_filename(cell_id: CellId_t, suffix: str = "") -> str:
     return os.path.join(get_tmpdir(), basename + suffix + ".py")
 
 
+def ends_with_semicolon(code: str) -> bool:
+    """Returns True if the cell's code ends with a semicolon, ignoring whitespace and comments.
+
+    Args:
+        code: The cell's source code
+    Returns:
+        bool: True if the last non-comment line ends with a semicolon
+    """
+    # Tokenize to check for semicolon
+    tokens = tokenize(io.BytesIO(code.strip().encode("utf-8")).readline)
+    for token in reversed(list(tokens)):
+        if token.type in (
+            token_types.ENDMARKER,
+            token_types.NEWLINE,
+            token_types.NL,
+            token_types.COMMENT,
+        ):
+            continue
+        return token.string == ";"
+    return False
+
+
 def cache(filename: str, code: str) -> None:
     # Generate a cache entry in Python's linecache
     linecache.cache[filename] = (
@@ -139,7 +161,9 @@ def compile_cell(
 
     expr: ast.Expression
     final_expr = module.body[-1]
-    if isinstance(final_expr, ast.Expr):
+    # Use final expression if it exists doesn't end in a
+    # semicolon. Evaluates expression to "None" otherwise.
+    if isinstance(final_expr, ast.Expr) and not ends_with_semicolon(code):
         expr = ast.Expression(module.body.pop().value)
         expr.lineno = final_expr.lineno
     else:
diff --git a/tests/_ast/test_compiler.py b/tests/_ast/test_compiler.py
index 0add58d903a..c2be07b5baa 100644
--- a/tests/_ast/test_compiler.py
+++ b/tests/_ast/test_compiler.py
@@ -347,3 +347,83 @@ def test_cell_id_from_filename() -> None:
     )
 
     assert compiler.cell_id_from_filename("random_file.py") is None
+
+
+class TestSemicolon:
+    @staticmethod
+    def test_return() -> None:
+        # fmt: off
+        def f() -> None:
+            1  # noqa: B018
+        # fmt: on
+
+        cell = compiler.cell_factory(f, cell_id="0")
+        assert eval(cell._cell.last_expr) == 1
+
+    @staticmethod
+    def test_return_suppressed() -> None:
+        # fmt: off
+        def f() -> None:
+            1;  # noqa: B018 E703
+        # fmt: on
+
+        cell = compiler.cell_factory(f, cell_id="0")
+        assert eval(cell._cell.last_expr) is None
+
+    @staticmethod
+    def test_return_last() -> None:
+        # fmt: off
+        def f() -> None:
+            1; 2; 3  # noqa: B018 E702
+        # fmt: on
+
+        cell = compiler.cell_factory(f, cell_id="0")
+        assert eval(cell._cell.last_expr) == 3
+
+    @staticmethod
+    def test_return_last_suppressed() -> None:
+        # fmt: off
+        def f() -> None:
+            1; 2; 3;  # noqa: B018 E702 E703
+        # fmt: on
+
+        cell = compiler.cell_factory(f, cell_id="0")
+        assert eval(cell._cell.last_expr) is None
+
+    @staticmethod
+    def test_return_comment() -> None:
+        def f() -> None:
+            1  # noqa: B018 # Has a comment;
+
+        cell = compiler.cell_factory(f, cell_id="0")
+        assert eval(cell._cell.last_expr) == 1
+
+    @staticmethod
+    def test_return_comment_suppressed() -> None:
+        # fmt: off
+        def f() -> None:
+            1;  # noqa: B018 E703 # Has a comment
+        # fmt: on
+
+        cell = compiler.cell_factory(f, cell_id="0")
+        assert eval(cell._cell.last_expr) is None
+
+    @staticmethod
+    def test_return_string_semicolon() -> None:
+        def f() -> None:
+            "#; splits on ;# are less than ideal"  # noqa: B018 Contains a ;#
+
+        cell = compiler.cell_factory(f, cell_id="0")
+        assert (
+            eval(cell._cell.last_expr) == "#; splits on ;# are less than ideal"
+        )
+
+    @staticmethod
+    def test_return_string_semicolon_suppressed() -> None:
+        # fmt: off
+        def f() -> None:
+            "#; splits on ;# are less than ideal";  # noqa: B018 E703 Contains a ;#
+        # fmt: on
+
+        cell = compiler.cell_factory(f, cell_id="0")
+        assert eval(cell._cell.last_expr) is None