diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000..dc3064b --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,14 @@ +{ + "version": "0.2.0", + "configurations": [ + { + "name": "Extension", + "type": "extensionHost", + "request": "launch", + "runtimeExecutable": "${execPath}", + "args": [ + "--extensionDevelopmentPath=${workspaceFolder}" + ] + } + ] +} \ No newline at end of file diff --git a/extension.js b/extension.js index deb44f5..5b3ba5a 100644 --- a/extension.js +++ b/extension.js @@ -1,8 +1,90 @@ const vscode = require('vscode'); const path = require('path'); +const { spawn } = require('child_process'); let sharedTerminal = null; +/** + * Tree Data Provider for Metaflow Steps + */ +class MetaflowTreeProvider { + constructor() { + this._onDidChangeTreeData = new vscode.EventEmitter(); + this.onDidChangeTreeData = this._onDidChangeTreeData.event; + this.steps = []; + } + + refresh() { + const editor = vscode.window.activeTextEditor; + if (!editor || !editor.document.fileName.endsWith('.py')) { + this.steps = []; + this._onDidChangeTreeData.fire(); + return; + } + + const pythonPath = vscode.workspace.getConfiguration('python').get('defaultInterpreterPath') || 'python'; + const parserPath = path.join(__dirname, 'parser.py'); + const filePath = editor.document.fileName; + + const pyParser = spawn(pythonPath, [parserPath, filePath]); + let output = ''; + + pyParser.stdout.on('data', (data) => { + output += data.toString(); + }); + + pyParser.stderr.on('data', (data) => { + console.error(`Parser error: ${data}`); + }); + + pyParser.on('close', (code) => { + if (code === 0) { + try { + this.steps = JSON.parse(output); + this._onDidChangeTreeData.fire(); + } catch (e) { + console.error('Failed to parse JSON output from Python script', e); + } + } + }); + } + + getTreeItem(element) { + const treeItem = new vscode.TreeItem(element.name); + treeItem.description = element.type; + treeItem.tooltip = `${element.name} (${element.type})`; + treeItem.command = { + command: 'metaflow.goToStep', + title: 'Go to Step', + arguments: [element.line], + }; + + // Set icons based on step type + switch (element.type) { + case 'linear': + treeItem.iconPath = new vscode.ThemeIcon('arrow-right'); + break; + case 'split': + treeItem.iconPath = new vscode.ThemeIcon('git-branch'); + break; + case 'foreach': + treeItem.iconPath = new vscode.ThemeIcon('repo-forked'); + break; + case 'join': + treeItem.iconPath = new vscode.ThemeIcon('git-merge'); + break; + default: + treeItem.iconPath = new vscode.ThemeIcon('circle-outline'); + } + + return treeItem; + } + + getChildren() { + return Promise.resolve(this.steps); + } +} + /** * Core function to detect the current Python function name and run a script. */ @@ -48,15 +130,12 @@ async function runPythonCommand(scriptName) { sharedTerminal.show(); sharedTerminal.sendText(`cd "${fileDir}"`); sharedTerminal.sendText(command); - - /* - vscode.window.showInformationMessage( - `${scriptName.toUpperCase()}: ${funcName} from ${path.basename(filePath)}` - ); - */ } function activate(context) { + const treeProvider = new MetaflowTreeProvider(); + vscode.window.registerTreeDataProvider('metaflow.stepOutline', treeProvider); + const runCmd = vscode.commands.registerCommand( 'extension.runPythonFunction', () => runPythonCommand('run_func') @@ -67,7 +146,45 @@ function activate(context) { () => runPythonCommand('spin_func') ); - context.subscriptions.push(runCmd, spinCmd); + const goToStepCmd = vscode.commands.registerCommand( + 'metaflow.goToStep', + (line) => { + const editor = vscode.window.activeTextEditor; + if (editor) { + const position = new vscode.Position(line - 1, 0); + editor.selection = new vscode.Selection(position, position); + editor.revealRange(new vscode.Range(position, position), vscode.TextEditorRevealType.InCenter); + } + } + ); + + const refreshCmd = vscode.commands.registerCommand('metaflow.refreshEntry', () => treeProvider.refresh()); + + // Auto-refresh logic + context.subscriptions.push( + vscode.window.onDidChangeActiveTextEditor(() => treeProvider.refresh()), + vscode.workspace.onDidSaveTextDocument((doc) => { + if (doc.fileName.endsWith('.py')) { + treeProvider.refresh(); + } + }) + ); + + // Debounced refresh on text change (optional but good) + let debounceTimer; + context.subscriptions.push( + vscode.workspace.onDidChangeTextDocument((e) => { + if (e.document === vscode.window.activeTextEditor?.document && e.document.fileName.endsWith('.py')) { + clearTimeout(debounceTimer); + debounceTimer = setTimeout(() => treeProvider.refresh(), 1000); + } + }) + ); + + context.subscriptions.push(runCmd, spinCmd, goToStepCmd, refreshCmd); + + // Initial refresh + treeProvider.refresh(); } function deactivate() { diff --git a/package.json b/package.json index 31a7c2f..11d71e5 100644 --- a/package.json +++ b/package.json @@ -1,40 +1,84 @@ { "name": "metaflow-dev", - "displayName": "Run and Spin Metaflow", - "version": "0.0.7", - "publisher": "local", - "engines": { "vscode": "^1.80.0" }, - "activationEvents": [ - "onCommand:extension.runPythonFunction", - "onCommand:extension.spinPythonFunction" - ], - "repository": { - "type": "git", - "url": "https://github.com/outerbounds/metaflow-dev-vscode" + "displayName": "Metaflow Dev Tools", + "description": "Shortcuts and tools for Metaflow development in VS Code", + "version": "0.0.8", + "publisher": "nandinisaagar", + "engines": { + "vscode": "^1.75.0" }, + "categories": [ + "Other" + ], + "activationEvents": [], "main": "./extension.js", "contributes": { + "viewsContainers": { + "activitybar": [ + { + "id": "metaflow-explorer", + "title": "Metaflow", + "icon": "$(rocket)" + } + ] + }, + "views": { + "metaflow-explorer": [ + { + "id": "metaflow.stepOutline", + "name": "Step Outline" + } + ] + }, "commands": [ { "command": "extension.runPythonFunction", - "title": "Run a flow" + "title": "Metaflow: Run Current Step" }, { "command": "extension.spinPythonFunction", - "title": "Spin the current step" + "title": "Metaflow: Spin Current Step" + }, + { + "command": "metaflow.goToStep", + "title": "Metaflow: Go to Step" + }, + { + "command": "metaflow.refreshEntry", + "title": "Refresh", + "icon": "$(refresh)" } ], + "menus": { + "view/title": [ + { + "command": "metaflow.refreshEntry", + "when": "view == metaflow.stepOutline", + "group": "navigation" + } + ] + }, "keybindings": [ { "command": "extension.runPythonFunction", "key": "ctrl+alt+r", - "when": "editorLangId == python" + "mac": "cmd+alt+r" }, { "command": "extension.spinPythonFunction", "key": "ctrl+alt+s", - "when": "editorLangId == python" + "mac": "cmd+alt+s" } ] + }, + "scripts": { + "lint": "eslint .", + "pretest": "npm run lint", + "test": "node ./test/runTest.js" + }, + "devDependencies": { + "@types/vscode": "^1.75.0", + "@types/node": "16.x", + "eslint": "^8.34.0" } } diff --git a/parser.py b/parser.py new file mode 100644 index 0000000..c3c9336 --- /dev/null +++ b/parser.py @@ -0,0 +1,139 @@ +import ast +import json +import sys +import os + +class MetaflowParser(ast.NodeVisitor): + def __init__(self): + self.steps = {} + self.flow_class = None + + def visit_ClassDef(self, node): + # Check if inherits from FlowSpec + for base in node.bases: + if isinstance(base, ast.Name) and base.id == 'FlowSpec': + self.flow_class = node.name + break + elif isinstance(base, ast.Attribute) and base.attr == 'FlowSpec': + self.flow_class = node.name + break + + if self.flow_class: + self.generic_visit(node) + + def visit_FunctionDef(self, node): + # Check for @step decorator + is_step = False + for decorator in node.decorator_list: + if isinstance(decorator, ast.Name) and decorator.id == 'step': + is_step = True + break + elif isinstance(decorator, ast.Attribute) and decorator.attr == 'step': + is_step = True + break + + if is_step: + step_name = node.name + next_steps = [] + step_type = 'linear' + + # Look for self.next(...) + for subnode in ast.walk(node): + if isinstance(subnode, ast.Call): + if isinstance(subnode.func, ast.Attribute) and \ + isinstance(subnode.func.value, ast.Name) and \ + subnode.func.value.id == 'self' and \ + subnode.func.attr == 'next': + + # Extract arguments + args = [] + for arg in subnode.args: + if isinstance(arg, ast.Attribute) and \ + isinstance(arg.value, ast.Name) and \ + arg.value.id == 'self': + args.append(arg.attr) + + # Extract foreach + is_foreach = False + for keyword in subnode.keywords: + if keyword.arg == 'foreach': + is_foreach = True + break + + if is_foreach: + step_type = 'foreach' + elif len(args) > 1: + step_type = 'split' + + next_steps.extend(args) + + # Check if it's a join step (accepts 'inputs' argument) + for arg in node.args.args: + if arg.arg == 'inputs': + step_type = 'join' + break + + self.steps[step_name] = { + 'name': step_name, + 'line': node.lineno, + 'type': step_type, + 'next': next_steps + } + +def get_dag_order(steps): + if not steps: + return [] + + # Simple Topological Sort/BFS from 'start' + ordered = [] + visited = set() + queue = ['start'] if 'start' in steps else [] + + # If no 'start', just return alphabetically but 'start' is usually there + if not queue and steps: + queue = [sorted(steps.keys())[0]] + + while queue: + curr = queue.pop(0) + if curr in visited or curr not in steps: + continue + + visited.add(curr) + ordered.append(steps[curr]) + + for nxt in steps[curr]['next']: + if nxt not in visited: + queue.append(nxt) + + # Add any remaining steps that weren't reachable (e.g. disconnected components) + for step_name in sorted(steps.keys()): + if step_name not in visited: + ordered.append(steps[step_name]) + + return ordered + +def main(): + if len(sys.argv) < 2: + print(json.dumps({"error": "No file path provided"})) + return + + file_path = sys.argv[1] + if not os.path.exists(file_path): + print(json.dumps({"error": f"File not found: {file_path}"})) + return + + try: + with open(file_path, 'r') as f: + tree = ast.parse(f.read()) + + parser = MetaflowParser() + parser.visit(tree) + + dag_ordered_steps = get_dag_order(parser.steps) + print(json.dumps(dag_ordered_steps)) + + except Exception as e: + print(json.dumps({"error": str(e)})) + +if __name__ == "__main__": + main() diff --git a/sample_flow.py b/sample_flow.py new file mode 100644 index 0000000..c5a6437 --- /dev/null +++ b/sample_flow.py @@ -0,0 +1,58 @@ +from metaflow import FlowSpec, step + +class SampleFlow(FlowSpec): + @step + def start(self): + print("Starting...") + self.next(self.intermediate_step) + + @step + def intermediate_step(self): + print("Intermediate step...") + self.next(self.linear_step) + + @step + def linear_step(self): + print("Linear step") + self.next(self.split_step) + + @step + def split_step(self): + print("Splitting...") + self.next(self.branch_a, self.branch_b) + + @step + def branch_a(self): + print("Branch A") + self.next(self.join_step) + + @step + def branch_b(self): + print("Branch B") + self.next(self.join_step) + + @step + def join_step(self, inputs): + print("Joining...") + self.next(self.foreach_step) + + @step + def foreach_step(self): + self.items = [1, 2, 3] + self.next(self.analyze_item, foreach='items') + + @step + def analyze_item(self): + print("Analyzing item", self.input) + self.next(self.end_join) + + @step + def end_join(self, inputs): + self.next(self.end) + + @step + def end(self): + print("End.") + +if __name__ == '__main__': + SampleFlow()