add flow analyser for diagnostics

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() {

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<string, Set<string>>} */
+    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,
+};
+

package.json

@@ -5,6 +5,7 @@
   "publisher": "local",
   "engines": { "vscode": "^1.80.0" },
   "activationEvents": [
+    "onLanguage:python",
     "onCommand:extension.runPythonFunction",
     "onCommand:extension.spinPythonFunction"
   ],