SONARPY-3984 Fix S3626 false positive for bare return in marimo notebook cells (#1023)

Co-authored-by: Vibe Bot <vibe-bot@sonarsource.com>
GitOrigin-RevId: ac09baa64a490f05145a060ebca1d205ca8ef455

Author: 2 people

python-checks/src/main/java/org/sonar/python/checks/RedundantJumpCheck.java

@@ -31,6 +31,7 @@
 import org.sonar.plugins.python.api.tree.Tree.Kind;
 import org.sonar.python.cfg.PythonCfgBranchingBlock;
 import org.sonar.plugins.python.api.quickfix.PythonQuickFix;
+import org.sonar.python.checks.utils.MarimoUtils;
 import org.sonar.python.quickfix.TextEditUtils;
 import org.sonar.python.tree.TreeUtils;
 
@@ -54,7 +55,7 @@ private static void checkCfg(@Nullable ControlFlowGraph cfg, SubscriptionContext
       .forEach(cfgBlock -> {
         List<Tree> elements = cfgBlock.elements();
         Tree lastElement = elements.get(elements.size() - 1);
-        if (!isException(lastElement)) {
+        if (!isException(lastElement, ctx)) {
           var issue = ctx.addIssue(lastElement, message(lastElement));
           addQuickFix(lastElement, issue);
 
@@ -92,11 +93,12 @@ private static boolean isReturnWithExpression(Tree lastElement) {
     return lastElement.is(Kind.RETURN_STMT) && !((ReturnStatement) lastElement).expressions().isEmpty();
   }
 
-  private static boolean isException(Tree lastElement) {
+  private static boolean isException(Tree lastElement, SubscriptionContext ctx) {
     return lastElement.is(Kind.RAISE_STMT)
       || isReturnWithExpression(lastElement)
       || isInsideSingleStatementBlock(lastElement)
-      || hasTryAncestor(lastElement);
+      || hasTryAncestor(lastElement)
+      || MarimoUtils.isTreeInMarimoDecoratedFunction(lastElement, ctx);
   }
 
   // ignore jumps in try statement because CFG is not precise

python-checks/src/main/java/org/sonar/python/checks/utils/MarimoUtils.java

@@ -0,0 +1,53 @@
+/*
+ * SonarQube Python Plugin
+ * Copyright (C) 2011-2025 SonarSource Sàrl
+ * mailto:info AT sonarsource DOT com
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the Sonar Source-Available License Version 1, as published by SonarSource SA.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+ * See the Sonar Source-Available License for more details.
+ *
+ * You should have received a copy of the Sonar Source-Available License
+ * along with this program; if not, see https://sonarsource.com/license/ssal/
+ */
+package org.sonar.python.checks.utils;
+
+import org.sonar.plugins.python.api.SubscriptionContext;
+import org.sonar.plugins.python.api.tree.CallExpression;
+import org.sonar.plugins.python.api.tree.Decorator;
+import org.sonar.plugins.python.api.tree.Expression;
+import org.sonar.plugins.python.api.tree.FunctionDef;
+import org.sonar.plugins.python.api.tree.Tree;
+import org.sonar.plugins.python.api.tree.Tree.Kind;
+import org.sonar.plugins.python.api.types.v2.matchers.TypeMatcher;
+import org.sonar.plugins.python.api.types.v2.matchers.TypeMatchers;
+import org.sonar.python.tree.TreeUtils;
+
+public class MarimoUtils {
+
+  private static final TypeMatcher APP_CELL_MATCHER = TypeMatchers.isType("marimo.App.cell");
+  private static final TypeMatcher APP_FUNCTION_MATCHER = TypeMatchers.isType("marimo.App.function");
+
+  private MarimoUtils() {
+  }
+
+  public static boolean isTreeInMarimoDecoratedFunction(Tree tree, SubscriptionContext ctx) {
+    FunctionDef functionDef = (FunctionDef) TreeUtils.firstAncestorOfKind(tree, Kind.FUNCDEF);
+    if (functionDef == null) {
+      return false;
+    }
+    return functionDef.decorators().stream().anyMatch(decorator -> isAppCellDecorator(decorator, ctx));
+  }
+
+  private static boolean isAppCellDecorator(Decorator decorator, SubscriptionContext ctx) {
+    Expression expr = decorator.expression();
+    if (expr instanceof CallExpression callExpr) {
+      expr = callExpr.callee();
+    }
+    return APP_CELL_MATCHER.isTrueFor(expr, ctx) || APP_FUNCTION_MATCHER.isTrueFor(expr, ctx);
+  }
+}

python-checks/src/test/java/org/sonar/python/checks/utils/MarimoUtilsTest.java

@@ -0,0 +1,86 @@
+/*
+ * SonarQube Python Plugin
+ * Copyright (C) 2011-2025 SonarSource Sàrl
+ * mailto:info AT sonarsource DOT com
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the Sonar Source-Available License Version 1, as published by SonarSource SA.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+ * See the Sonar Source-Available License for more details.
+ *
+ * You should have received a copy of the Sonar Source-Available License
+ * along with this program; if not, see https://sonarsource.com/license/ssal/
+ */
+package org.sonar.python.checks.utils;
+
+import com.sonar.sslr.api.AstNode;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+import org.sonar.plugins.python.api.SubscriptionContext;
+import org.sonar.plugins.python.api.tree.FileInput;
+import org.sonar.plugins.python.api.tree.FunctionDef;
+import org.sonar.python.TestPythonVisitorRunner;
+import org.sonar.python.parser.PythonParser;
+import org.sonar.python.semantic.ProjectLevelSymbolTable;
+import org.sonar.python.semantic.v2.SymbolTableBuilderV2;
+import org.sonar.python.semantic.v2.TypeInferenceV2;
+import org.sonar.python.semantic.v2.typetable.ProjectLevelTypeTable;
+import org.sonar.python.tree.PythonTreeMaker;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class MarimoUtilsTest {
+
+  @Test
+  void isTreeInMarimoDecoratedFunction_returnsFalseWhenNotInsideFunction() {
+    FileInput fileInput = parse("""
+      import marimo
+      x = 1
+      """);
+
+    SubscriptionContext ctx = mockSubscriptionContext();
+
+    assertThat(MarimoUtils.isTreeInMarimoDecoratedFunction(fileInput, ctx)).isFalse();
+  }
+
+  @Test
+  void isTreeInMarimoDecoratedFunction_returnsTrueWhenInsideAppCellFunction() {
+    String content = """
+      import marimo
+      app = marimo.App()
+      @app.cell
+      def some_cell():
+          return 1
+      """;
+
+    FileInput fileInput = parse(content);
+    var typeTable = new ProjectLevelTypeTable(ProjectLevelSymbolTable.empty());
+    var mockFile = new TestPythonVisitorRunner.MockPythonFile("", "test.py", content);
+    var symbolTableV2 = new SymbolTableBuilderV2(fileInput).build();
+    new TypeInferenceV2(typeTable, mockFile, symbolTableV2, "").inferModuleType(fileInput);
+
+    SubscriptionContext ctx = Mockito.mock(SubscriptionContext.class);
+    Mockito.when(ctx.typeTable()).thenReturn(typeTable);
+
+    var statements = fileInput.statements().statements();
+    var functionDef = (FunctionDef) statements.get(statements.size() - 1);
+    var returnStatement = functionDef.body().statements().get(0);
+
+    assertThat(MarimoUtils.isTreeInMarimoDecoratedFunction(returnStatement, ctx)).isTrue();
+  }
+
+  private static SubscriptionContext mockSubscriptionContext() {
+    SubscriptionContext ctx = Mockito.mock(SubscriptionContext.class);
+    Mockito.when(ctx.typeTable()).thenReturn(new ProjectLevelTypeTable(ProjectLevelSymbolTable.empty()));
+    return ctx;
+  }
+
+  private static FileInput parse(String content) {
+    PythonParser parser = PythonParser.create();
+    AstNode astNode = parser.parse(content);
+    return new PythonTreeMaker().fileInput(astNode);
+  }
+}

python-checks/src/test/resources/checks/redundantJump.py

@@ -340,3 +340,47 @@ def match_statement_single_return(value):
     match value:
       case "42":
         return  # OK
+
+
+import marimo
+
+app = marimo.App()
+
+
+@app.cell
+def _(mo, engine):
+    _df = mo.sql(f"""SELECT * FROM table""", engine=engine)
+    return  # OK
+
+
+@app.cell(hide_code=True)
+def _(mo):
+    mo.md("# Dashboard")
+    return  # OK
+
+
+@app.cell
+def _(mo, condition):
+    if condition:
+        mo.md("Condition met")
+        return  # OK
+    mo.md("Condition not met")
+    return  # OK
+
+
+@app.function
+def helper():
+    x = 1
+    return  # OK
+
+
+@app.function
+def helper_with_condition(condition):
+    if condition:
+        return  # OK
+    return  # OK
+
+
+def not_a_marimo_cell():
+    x = 1
+    return  # Noncompliant

python-frontend/src/main/resources/org/sonar/python/types/custom_protobuf/marimo.protobuf

@@ -0,0 +1,41 @@
+
+marimo�
+App
+marimo.App"*SonarPythonAnalyzerFakeStub.CustomStubBase*u
+__init__marimo.App.__init__"
+None*"
+self

+
+marimo.App"
+marimo.App*
+args
+Any*
+kwargs
+Any*l
+cellmarimo.App.cell"
+Any*"
+self

+
+marimo.App"
+marimo.App*
+args
+Any*
+kwargs
+Any*t
+functionmarimo.App.function"
+Any*"
+self

+
+marimo.App"
+marimo.App*
+args
+Any*
+kwargs
+Any*g
+__path__marimo.__path__J
+builtins.list[builtins.str]
+builtins.str"builtins.str"builtins.list*�

+__annotations__marimo.__annotations__W
+builtins.dict[builtins.str,Any]
+builtins.str"builtins.str
+Any"builtins.dict

python-frontend/typeshed_serializer/resources/custom/marimo/__init__.pyi

@@ -0,0 +1,7 @@
+from SonarPythonAnalyzerFakeStub import CustomStubBase
+from typing import Any
+
+class App(CustomStubBase):
+    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
+    def cell(self, *args: Any, **kwargs: Any) -> Any: ...
+    def function(self, *args: Any, **kwargs: Any) -> Any: ...

python-frontend/typeshed_serializer/tests/test_serializers.py

@@ -74,7 +74,7 @@ def test_custom_stubs_serializer(typeshed_custom_stubs):
         custom_stubs_serializer.serialize()
         assert custom_stubs_serializer.get_build_result.call_count == 1
         # Not every files from "typeshed_custom_stubs" build are serialized, as some are builtins
-        assert symbols.save_module.call_count == 361
+        assert symbols.save_module.call_count == 362
 
 
 def test_importer_serializer():