diff --git a/extension.js b/extension.js index deb44f5..05c1029 100644 --- a/extension.js +++ b/extension.js @@ -1,5 +1,6 @@ const vscode = require('vscode'); const path = require('path'); +const { analyzeMetaflowFlow } = require('./metaflowDiagnostics'); let sharedTerminal = null; @@ -56,6 +57,38 @@ async function runPythonCommand(scriptName) { */ } +function refreshMetaflowDiagnostics(document, collection) { + if (!document || document.languageId !== 'python') { + return; + } + + const text = document.getText(); + const rawDiagnostics = analyzeMetaflowFlow(text); + + const diagnostics = rawDiagnostics.map((d) => { + const range = new vscode.Range( + new vscode.Position(d.startLine, d.startChar), + new vscode.Position(d.endLine, d.endChar) + ); + + let severity = vscode.DiagnosticSeverity.Warning; + if (d.severity === 'error') { + severity = vscode.DiagnosticSeverity.Error; + } else if (d.severity === 'info') { + severity = vscode.DiagnosticSeverity.Information; + } + + const diag = new vscode.Diagnostic(range, d.message, severity); + diag.source = 'metaflow-dev'; + if (d.code) { + diag.code = d.code; + } + return diag; + }); + + collection.set(document.uri, diagnostics); +} + function activate(context) { const runCmd = vscode.commands.registerCommand( 'extension.runPythonFunction', @@ -67,7 +100,31 @@ function activate(context) { () => runPythonCommand('spin_func') ); - context.subscriptions.push(runCmd, spinCmd); + const diagnosticCollection = vscode.languages.createDiagnosticCollection( + 'metaflow-flow' + ); + + context.subscriptions.push(runCmd, spinCmd, diagnosticCollection); + + context.subscriptions.push( + vscode.workspace.onDidOpenTextDocument((doc) => + refreshMetaflowDiagnostics(doc, diagnosticCollection) + ), + vscode.workspace.onDidChangeTextDocument((event) => + refreshMetaflowDiagnostics(event.document, diagnosticCollection) + ), + vscode.workspace.onDidSaveTextDocument((doc) => + refreshMetaflowDiagnostics(doc, diagnosticCollection) + ), + vscode.workspace.onDidCloseTextDocument((doc) => + diagnosticCollection.delete(doc.uri) + ) + ); + + const editor = vscode.window.activeTextEditor; + if (editor && editor.document.languageId === 'python') { + refreshMetaflowDiagnostics(editor.document, diagnosticCollection); + } } function deactivate() { diff --git a/metaflowDiagnostics.js b/metaflowDiagnostics.js new file mode 100644 index 0000000..3187605 --- /dev/null +++ b/metaflowDiagnostics.js @@ -0,0 +1,276 @@ +function analyzeMetaflowFlow(text) { + const lines = text.split(/\r?\n/); + + const flowClasses = []; + + let currentClass = null; + let pendingDecorators = []; + + const getIndent = (line) => { + const m = line.match(/^(\s*)/); + return m ? m[1].length : 0; + }; + + for (let i = 0; i < lines.length; i++) { + const line = lines[i]; + const trimmed = line.trim(); + const indent = getIndent(line); + + if ( + currentClass && + indent <= currentClass.indent && + trimmed !== '' && + !trimmed.startsWith('#') && + !trimmed.startsWith('@') && + !trimmed.startsWith('def ') + ) { + currentClass = null; + pendingDecorators = []; + } + + if (trimmed.startsWith('class ')) { + const classMatch = trimmed.match( + /^class\s+([A-Za-z_][A-Za-z0-9_]*)\s*\(([^)]*)\)\s*:/ + ); + if (classMatch) { + const baseList = classMatch[2] || ''; + const isFlowSpec = /\bFlowSpec\b/.test(baseList); + if (isFlowSpec) { + currentClass = { + name: classMatch[1], + line: i, + indent, + steps: [], + edges: [], + }; + flowClasses.push(currentClass); + pendingDecorators = []; + continue; + } + } + currentClass = null; + pendingDecorators = []; + continue; + } + + if (!currentClass) { + pendingDecorators = []; + continue; + } + + if (trimmed.startsWith('@')) { + pendingDecorators.push({ lineIndex: i, text: trimmed }); + continue; + } + + const defMatch = trimmed.match(/^def\s+([A-Za-z_][A-Za-z0-9_]*)\s*\(/); + if (defMatch && indent > currentClass.indent) { + const hasStepDecorator = pendingDecorators.some((d) => + /^@step\b/.test(d.text) + ); + pendingDecorators = []; + + if (hasStepDecorator) { + const name = defMatch[1]; + const nameStartChar = line.indexOf(name); + const step = { + name, + line: i, + nameStartChar, + nameEndChar: nameStartChar + name.length, + indent, + bodyStartLine: null, + bodyEndLine: null, + }; + currentClass.steps.push(step); + } + } else if (trimmed !== '' && !trimmed.startsWith('#')) { + pendingDecorators = []; + } + } + + /** @type {MetaflowDiagnostic[]} */ + const diagnostics = []; + + if (flowClasses.length === 0) { + return diagnostics; + } + + for (const cls of flowClasses) { + for (const step of cls.steps) { + let bodyStart = null; + let j = step.line + 1; + while (j < lines.length) { + const line = lines[j]; + const trimmed = line.trim(); + const indent = getIndent(line); + + if (trimmed === '' || trimmed.startsWith('#')) { + j += 1; + continue; + } + + if (indent <= step.indent) { + break; + } + + bodyStart = j; + j += 1; + break; + } + + if (bodyStart === null) { + step.bodyStartLine = null; + step.bodyEndLine = null; + continue; + } + + let bodyEnd = bodyStart; + let k = bodyStart; + while (k < lines.length) { + const line = lines[k]; + const trimmed = line.trim(); + const indent = getIndent(line); + if (trimmed !== '' && !trimmed.startsWith('#') && indent <= step.indent) { + break; + } + bodyEnd = k + 1; + k += 1; + } + + step.bodyStartLine = bodyStart; + step.bodyEndLine = bodyEnd; + } + } + + for (const cls of flowClasses) { + const stepNames = new Set(cls.steps.map((s) => s.name)); + /** @type {Record>} */ + const edgesFrom = {}; + const hasNextCall = {}; + + for (const step of cls.steps) { + edgesFrom[step.name] = new Set(); + hasNextCall[step.name] = false; + + if (step.bodyStartLine == null || step.bodyEndLine == null) { + continue; + } + + for (let i = step.bodyStartLine; i < step.bodyEndLine; i++) { + const lineText = lines[i]; + const nextRegex = /self\.next\s*\(([^)]*)\)/g; + let nextMatch; + + while ((nextMatch = nextRegex.exec(lineText)) !== null) { + hasNextCall[step.name] = true; + + const args = nextMatch[1] || ''; + const nameRegex = /self\.([A-Za-z_][A-Za-z0-9_]*)/g; + let nameMatch; + while ((nameMatch = nameRegex.exec(args)) !== null) { + const targetName = nameMatch[1]; + const nameIndex = lineText.indexOf( + targetName, + nextMatch.index + ); + const startChar = + nameIndex >= 0 ? nameIndex : nextMatch.index; + const endChar = startChar + targetName.length; + + if (stepNames.has(targetName)) { + edgesFrom[step.name].add(targetName); + } else { + diagnostics.push({ + message: `self.next refers to undefined step method '${targetName}'.`, + severity: 'warning', + startLine: i, + startChar, + endLine: i, + endChar, + code: 'metaflow.undefinedStep', + }); + } + } + } + } + } + + for (const step of cls.steps) { + if (step.name === 'end') continue; + if (!hasNextCall[step.name]) { + diagnostics.push({ + message: + `Step '${step.name}' does not call self.next(); Metaflow steps should transition to another step (except 'end').`, + severity: 'warning', + startLine: step.line, + startChar: step.nameStartChar, + endLine: step.line, + endChar: step.nameEndChar, + code: 'metaflow.missingNext', + }); + } + } + + if (!stepNames.has('start')) { + diagnostics.push({ + message: `Metaflow flow '${cls.name}' is missing a 'start' step.`, + severity: 'warning', + startLine: cls.line, + startChar: 0, + endLine: cls.line, + endChar: lines[cls.line] ? lines[cls.line].length : 0, + code: 'metaflow.missingStart', + }); + } + + if (!stepNames.has('end')) { + diagnostics.push({ + message: `Metaflow flow '${cls.name}' is missing an 'end' step.`, + severity: 'warning', + startLine: cls.line, + startChar: 0, + endLine: cls.line, + endChar: lines[cls.line] ? lines[cls.line].length : 0, + code: 'metaflow.missingEnd', + }); + } + + if (stepNames.has('start')) { + const visited = new Set(); + const worklist = ['start']; + while (worklist.length) { + const current = worklist.pop(); + if (visited.has(current)) continue; + visited.add(current); + const neighbors = edgesFrom[current]; + if (!neighbors) continue; + for (const n of neighbors) { + if (!visited.has(n)) worklist.push(n); + } + } + + for (const step of cls.steps) { + if (step.name === 'start') continue; + if (!visited.has(step.name)) { + diagnostics.push({ + message: `Step '${step.name}' is not reachable from 'start'.`, + severity: 'warning', + startLine: step.line, + startChar: step.nameStartChar, + endLine: step.line, + endChar: step.nameEndChar, + code: 'metaflow.unreachableStep', + }); + } + } + } + } + + return diagnostics; +} + +module.exports = { + analyzeMetaflowFlow, +}; + diff --git a/package.json b/package.json index 31a7c2f..a84bb14 100644 --- a/package.json +++ b/package.json @@ -5,6 +5,7 @@ "publisher": "local", "engines": { "vscode": "^1.80.0" }, "activationEvents": [ + "onLanguage:python", "onCommand:extension.runPythonFunction", "onCommand:extension.spinPythonFunction" ],