diff --git a/src/pythonParsing.ts b/src/pythonParsing.ts index de43f76a..2c92b7ca 100644 --- a/src/pythonParsing.ts +++ b/src/pythonParsing.ts @@ -1,5 +1,5 @@ import * as crypto from "crypto"; -import { TextDocument, Range } from "vscode"; +import { Range, TextDocument } from "vscode"; /** * Cache is a simple key-value store that keeps a maximum number of entries. @@ -206,7 +206,7 @@ export class ManimClass { /** * Regular expression to match the construct() method definition. */ - private static CONSTRUCT_METHOD_REGEX = /^\s*def\s+construct\s*\(self\)\s*:/; + private static CONSTRUCT_METHOD_REGEX = /^\s*def\s+construct\s*\(self\)\s*(->\s*None)?\s*:/; /** * The 0-based line number where the Manim Class is defined. diff --git a/tests/cellRanges.test.ts b/tests/cellRanges.test.ts index 15629f59..1315f79b 100644 --- a/tests/cellRanges.test.ts +++ b/tests/cellRanges.test.ts @@ -1,7 +1,7 @@ -import { workspace, Range } from "vscode"; import { describe, it } from "mocha"; -import { uriInWorkspace } from "./utils/testRunner"; +import { Range, workspace } from "vscode"; import { ManimCellRanges } from "../src/pythonParsing"; +import { uriInWorkspace } from "./utils/testRunner"; describe("Manim Cell Ranges", function () { // in the expected ranges we only care about the start and end lines @@ -9,7 +9,7 @@ describe("Manim Cell Ranges", function () { const tests = [ { filename: "detection_basic.py", - expectedRanges: [[5, 7], [9, 10]], + expectedRanges: [[5, 7], [9, 10], [16, 18]], }, { filename: "detection_class_definition.py", diff --git a/tests/fixtures/detection_basic.py b/tests/fixtures/detection_basic.py index 60a64617..536b1f78 100644 --- a/tests/fixtures/detection_basic.py +++ b/tests/fixtures/detection_basic.py @@ -11,6 +11,14 @@ def construct(self): print("And even more code") +class BasicNotebookWithType(Scene): + + def construct(self) -> None: + ## Cell inside construct(self) marked with "None" type + print("With some code None") + print("With some more code None") + + class NoManimScene(Scene): def constructtttt(self): ## Should not be detected as Manim Cell