Skip to content

Add more tests for MarimoFileDetector and allow optional type annotation #138

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 16, 2025
Merged

Add more tests for MarimoFileDetector and allow optional type annotation #138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions extension/src/layers/MarimoFileDetector.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,9 @@ export const MarimoFileDetectorLive = Layer.scopedDiscard(

const text = document.getText();

// Check for top-level marimo.App() declaration (at start of line)
// This regex matches: optional whitespace, 'app', whitespace, '=', whitespace, 'marimo.App()'
return /^app\s*=\s*marimo\.App\(\)/m.test(text);
// Check for top-level marimo.App( declaration (at start of line)
// This regex matches: optional whitespace, 'app', optional type annotation, '=', whitespace, 'marimo.App()'
return /^app\s*(?::\s*[^=]+)?\s*=\s*marimo\.App\(/m.test(text);
};

// Update context based on active editor
Expand All @@ -42,7 +42,7 @@ export const MarimoFileDetectorLive = Layer.scopedDiscard(
yield* Effect.logDebug("Detected marimo notebook file").pipe(
Effect.annotateLogs({
uri: Option.map(editor, (e) => e.document.uri.toString()).pipe(
Option.getOrNull,
Option.getOrThrow,
),
}),
);
Expand Down
206 changes: 150 additions & 56 deletions extension/src/layers/__tests__/MarimoFileDetector.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,35 +22,78 @@ it.effect(
const ctx = yield* withTestCtx();
// build the layer
yield* Effect.provide(Effect.void, ctx.layer);
expect(yield* ctx.vscode.executions).toMatchInlineSnapshot(`
[
{
"args": [
"marimo.isPythonFileMarimoNotebook",
false,
],
"command": "setContext",
},
]
`);
expect(yield* ctx.vscode.executions).toEqual([
{
command: "setContext",
args: ["marimo.isPythonFileMarimoNotebook", false],
},
]);
}),
);

it.effect(
"should be true on intialization with active editor",
Effect.fnUntraced(function* () {
const ctx = yield* withTestCtx();
const pythonCode = `import marimo
it.effect.each([
[
"basic marimo app",
`import marimo

app = marimo.App()
`,
],
[
"marimo app with kwargs",
`import marimo

app = marimo.App(some=10, kwargs="20")
`,
],
["minimal marimo app without import", `app = marimo.App()`],
[
"marimo app with no whitespace around equals",
`import marimo
app=marimo.App()`,
],
[
"marimo app with extra whitespace",
`import marimo
app = marimo.App()`,
],
[
"marimo app with tabs",
`import marimo
app\t=\tmarimo.App()`,
],
[
"marimo app with comment above",
`import marimo

# Initialize the app
app = marimo.App()`,
],
[
"marimo app with type annotation",
`import marimo

app: marimo.App = marimo.App()`,
],
[
"complete marimo notebook with cells",
`import marimo

app = marimo.App(width="medium")

@app.cell
def __():
import numpy as np
return

if __name__ == "__main__":
app.run()
`;
`,
],
] as const)(
"should be true on intialization with active editor: %s",
Effect.fnUntraced(function* ([_, pythonCode]) {
const ctx = yield* withTestCtx();

const editor = createTestTextEditor(
createTestTextDocument("/test/notebook.py", "python", pythonCode),
Expand All @@ -60,17 +103,13 @@ if __name__ == "__main__":

// build the layer
yield* Effect.provide(Effect.void, ctx.layer);
expect(yield* ctx.vscode.executions).toMatchInlineSnapshot(`
[
{
"args": [
"marimo.isPythonFileMarimoNotebook",
true,
],
"command": "setContext",
},
]
`);

expect(yield* ctx.vscode.executions).toEqual([
{
command: "setContext",
args: ["marimo.isPythonFileMarimoNotebook", true],
},
]);
}),
);

Expand Down Expand Up @@ -104,31 +143,91 @@ if __name__ == "__main__":
yield* TestClock.adjust("100 millis");
}).pipe(Effect.provide(ctx.layer));

expect(yield* ctx.vscode.executions).toMatchInlineSnapshot(`
[
{
"args": [
"marimo.isPythonFileMarimoNotebook",
true,
],
"command": "setContext",
},
]
`);
expect(yield* ctx.vscode.executions).toEqual([
{
command: "setContext",
args: ["marimo.isPythonFileMarimoNotebook", true],
},
]);
}),
);

it.effect(
"should set context to false for regular Python file",
Effect.fnUntraced(function* () {
const ctx = yield* withTestCtx();
const pythonCode = `import pandas as pd
it.effect.each([
[
"regular Python script",
`import pandas as pd
def main():
print(pd)

if __name__ == "__main__":
main()
`;
`,
],
[
"imports marimo but doesn't create an app",
`import marimo

def helper():
pass
`,
],
[
"marimo.App in a string literal",
`code = "app = marimo.App()"
print(code)
`,
],
[
"marimo.App in a comment",
`# This file uses app = marimo.App() syntax
import other_lib
`,
],
[
"marimo.App indented inside a function",
`import marimo

def create_app():
app = marimo.App()
return app
`,
],
[
"marimo.App indented inside a class",
`import marimo

class NotebookManager:
def __init__(self):
self.app = marimo.App()
`,
],
[
"unrelated code mentioning marimo",
`my_app = some_other_framework.App()
marimo_reference = "check marimo.App docs"
`,
],
[
"marimo.App called without assignment",
`import marimo
marimo.App().run()
`,
],
[
"wrong variable name (not 'app')",
`import marimo
my_app = marimo.App()
`,
],
[
"only imports marimo",
`import marimo
`,
],
] as const)(
"should set context to false for non-marimo Python files: %s",
Effect.fnUntraced(function* ([_, pythonCode]) {
const ctx = yield* withTestCtx();
const editor = createTestTextEditor(
createTestTextDocument("/test/notebook.py", "python", pythonCode),
);
Expand All @@ -143,16 +242,11 @@ if __name__ == "__main__":
yield* TestClock.adjust("100 millis");
}).pipe(Effect.provide(ctx.layer));

expect(yield* ctx.vscode.executions).toMatchInlineSnapshot(`
[
{
"args": [
"marimo.isPythonFileMarimoNotebook",
false,
],
"command": "setContext",
},
]
`);
expect(yield* ctx.vscode.executions).toEqual([
{
command: "setContext",
args: ["marimo.isPythonFileMarimoNotebook", false],
},
]);
}),
);