diff --git a/README.md b/README.md index 02ff9ef..330debb 100644 --- a/README.md +++ b/README.md @@ -9,11 +9,37 @@ This lightweight VS Code extension supercharges your Metaflow dev workflow: The extension automatically: +* Validates that the active file contains a Metaflow `FlowSpec` class before running +* Detects `Parameter(...)` definitions and prompts you for values before execution +* Pre-fills default values and validates input types (int, float, bool, JSON) * Detects the current function name (`def` or `async def`) * Saves the file before running * Executes the command in the file’s directory * Reuses a shared terminal session +## Smart Parameter Prompts + +When you run a flow (Ctrl + Alt + R), the extension parses your Python file for `Parameter(...)` definitions and shows an input prompt for each one. Default values are pre-filled, required parameters are enforced, and type validation is applied automatically. + +The final command is built as `python flow.py run --param1=value1 --param2=value2` and sent to the terminal. + +### Current Limitations + +* Parameter detection uses static AST parsing — it does not import or execute your code. +* Dynamic parameters (created in loops or conditionals), renamed imports (e.g. `P = Parameter`), and deploy-time callable defaults are not detected. +* `IncludeFile` parameters are not yet supported. +* The extension uses `python` as the interpreter; custom virtualenv paths are not yet supported (see [#3](https://github.com/outerbounds/metaflow-dev-vscode/issues/3)). + +You can always fall back to running flows manually via the terminal if the prompt does not detect your parameters. + +## Running Tests + +```bash +npm test +``` + +This runs all tests using Node's built-in test runner. Tests cover flow detection, parameter extraction, and command construction. + ## Installation 1. **Clone or copy** this repository to a local folder diff --git a/buildRunCommand.js b/buildRunCommand.js new file mode 100644 index 0000000..b644c58 --- /dev/null +++ b/buildRunCommand.js @@ -0,0 +1,87 @@ +/** Metaflow/Click long-option style: letters, digits, underscore, hyphen (not whitespace or '='). */ +const CLI_NAME_RE = /^[a-zA-Z_][a-zA-Z0-9_-]*$/; + +/** + * Single-quote for bash/zsh: prevents all expansion. + * Only ' needs escaping, done by ending the string, adding an escaped + * literal quote, and reopening: it's → 'it'\''s' + */ +function quoteForBash(value) { + return "'" + value.replace(/'/g, "'\\''" ) + "'"; +} + +/** + * Single-quote for PowerShell: prevents all expansion. + * Only ' needs escaping, done by doubling: it's → 'it''s' + */ +function quoteForPowershell(value) { + return "'" + value.replace(/'/g, "''" ) + "'"; +} + +/** + * Double-quote for cmd.exe: the only option since cmd has no single-quote + * string semantics. Escape " by doubling and % by doubling. + */ +function quoteForCmd(value) { + return '"' + value.replace(/"/g, '""').replace(/%/g, '%%') + '"'; +} + +/** + * Quote a string for safe use in a terminal command. + * Uses single quotes for bash/PowerShell (prevents all variable expansion) + * and double quotes for cmd.exe. + * + * @param {string} value - The raw string to quote. + * @param {'bash'|'powershell'|'cmd'} shell - Target shell type. + */ +function quoteForTerminal(value, shell) { + if (shell === 'powershell') return quoteForPowershell(value); + if (shell === 'cmd') return quoteForCmd(value); + return quoteForBash(value); +} + +/** + * Detect shell type from a shell executable path. + * @param {string} shellPath - e.g. '/bin/bash' or 'C:\\...\\powershell.exe' + * @returns {'bash'|'powershell'|'cmd'} + */ +function detectShellType(shellPath) { + const lower = (shellPath || '').toLowerCase(); + if (lower.includes('powershell') || lower.includes('pwsh')) return 'powershell'; + if (lower.endsWith('cmd.exe') || lower.endsWith('cmd')) return 'cmd'; + return 'bash'; +} + +/** + * Build a `python run --flag=value ...` command string. + * + * @param {string} filePath - Absolute path to the flow file. + * @param {Array<{name: string, value: string}>} flagArgs - Parameter name/value pairs. + * @param {'bash'|'powershell'|'cmd'} shell - Target shell type. + * @returns {string} The shell command to execute. + */ +function buildRunCommand(filePath, flagArgs, shell) { + const parts = ['python', quoteForTerminal(filePath, shell), 'run']; + + for (const arg of flagArgs) { + if (!CLI_NAME_RE.test(arg.name)) { + throw new Error(`Invalid parameter name: ${arg.name}`); + } + if (/[\r\n]/.test(arg.value)) { + throw new Error(`Parameter value for --${arg.name} contains newline characters.`); + } + parts.push(`--${arg.name}=${quoteForTerminal(arg.value, shell)}`); + } + + return parts.join(' '); +} + +module.exports = { + buildRunCommand, + quoteForTerminal, + quoteForBash, + quoteForPowershell, + quoteForCmd, + detectShellType, + CLI_NAME_RE, +}; diff --git a/examples/hello_flow.py b/examples/hello_flow.py new file mode 100644 index 0000000..681be41 --- /dev/null +++ b/examples/hello_flow.py @@ -0,0 +1,36 @@ +""" +Example Metaflow flow for manual testing of the VS Code extension. + +To test parameter prompts: + 1. Open this file in VS Code with the extension installed + 2. Press Ctrl+Alt+R to run the flow + 3. You should see input prompts for 'greeting' and 'count' + pre-filled with their defaults + 4. The final terminal command should look like: + python hello_flow.py run --greeting='...' --count='...' + +To test spin: + 1. Place your cursor inside the 'start' or 'end' method + 2. Press Ctrl+Alt+S +""" + +from metaflow import FlowSpec, Parameter, step + + +class HelloFlow(FlowSpec): + greeting = Parameter('greeting', default='hello world', type=str, help='Message to print') + count = Parameter('count', default=3, type=int, help='How many times to repeat') + + @step + def start(self): + for i in range(self.count): + print(f"{i + 1}: {self.greeting}") + self.next(self.end) + + @step + def end(self): + print("Done!") + + +if __name__ == '__main__': + HelloFlow() diff --git a/extension.js b/extension.js index deb44f5..9bd0bbc 100644 --- a/extension.js +++ b/extension.js @@ -1,59 +1,100 @@ const vscode = require('vscode'); const path = require('path'); +const { validateEditorForFlowCommands, reportFlowInspectFailure } = require('./flowDetection'); +const { inspectFlowFile, promptForParameters } = require('./parameterPrompt'); +const { buildRunCommand, quoteForTerminal, detectShellType } = require('./buildRunCommand'); let sharedTerminal = null; -/** - * Core function to detect the current Python function name and run a script. - */ -async function runPythonCommand(scriptName) { - const editor = vscode.window.activeTextEditor; - if (!editor) return; +function getShellType() { + return detectShellType(vscode.env.shell); +} - const doc = editor.document; - await doc.save(); +function getOrCreateTerminal() { + if (!sharedTerminal || sharedTerminal.exitStatus !== undefined) { + sharedTerminal = vscode.window.createTerminal('Metaflow Runner'); + } + return sharedTerminal; +} - const cursorLine = editor.selection.active.line; +function sendToTerminal(fileDir, command) { + const shell = getShellType(); + const terminal = getOrCreateTerminal(); + terminal.show(); + terminal.sendText(`cd ${quoteForTerminal(fileDir, shell)}`); + terminal.sendText(command); +} - // Find the nearest "def"/"async def" above cursor - let funcName = null; - for (let i = cursorLine; i >= 0; i--) { +/** + * Find the nearest enclosing Python function name above the cursor. + */ +function findEnclosingFunction(doc, line) { + for (let i = line; i >= 0; i--) { const lineText = doc.lineAt(i).text.trim(); const match = lineText.match(/^(?:async\s+def|def)\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\(/); - if (match) { - funcName = match[1]; - break; - } + if (match) return match[1]; + } + return null; +} + +/** + * @param {import('vscode').TextDocument} doc + * @param {{ parameters: Array> }} inspectResult + */ +async function runFlow(doc, inspectResult) { + const filePath = doc.fileName; + const fileDir = path.dirname(filePath); + const shell = getShellType(); + + const params = inspectResult.parameters || []; + let flagArgs = []; + + if (params.length > 0) { + const values = await promptForParameters(params); + if (values === null) return; + flagArgs = values; } + try { + const command = buildRunCommand(filePath, flagArgs, shell); + sendToTerminal(fileDir, command); + } catch (err) { + console.error('Failed to build or run Metaflow command:', err); + vscode.window.showErrorMessage( + err instanceof Error ? err.message : 'Failed to run Metaflow flow.' + ); + } +} + +async function spinStep(doc, editor) { + const funcName = findEnclosingFunction(doc, editor.selection.active.line); if (!funcName) { - vscode.window.showErrorMessage("No enclosing Python function found."); + vscode.window.showErrorMessage('No enclosing Python function found.'); return; } const filePath = doc.fileName; const fileDir = path.dirname(filePath); + const shell = getShellType(); + const command = `python ${quoteForTerminal(filePath, shell)} spin ${funcName}`; + sendToTerminal(fileDir, command); +} - let command = ''; - if (scriptName == 'spin_func') - command = `python ${filePath} spin ${funcName}`; - else - command = `python ${filePath} run`; +async function runPythonCommand(scriptName) { + const editor = vscode.window.activeTextEditor; + const doc = validateEditorForFlowCommands(editor); + if (!doc) return; - // Reuse or create a single shared terminal - if (!sharedTerminal || sharedTerminal.exitStatus !== undefined) { - sharedTerminal = vscode.window.createTerminal('Metaflow Runner'); - } + await doc.save(); - sharedTerminal.show(); - sharedTerminal.sendText(`cd "${fileDir}"`); - sharedTerminal.sendText(command); + const inspect = await inspectFlowFile(doc.fileName); + if (reportFlowInspectFailure(inspect)) return; - /* - vscode.window.showInformationMessage( - `${scriptName.toUpperCase()}: ${funcName} from ${path.basename(filePath)}` - ); - */ + if (scriptName === 'spin_func') { + await spinStep(doc, editor); + } else { + await runFlow(doc, inspect); + } } function activate(context) { diff --git a/flowDetection.js b/flowDetection.js new file mode 100644 index 0000000..c3b4bb2 --- /dev/null +++ b/flowDetection.js @@ -0,0 +1,64 @@ +/** + * Pure predicate: whether AST inspection allows run/spin. + * @param {{ syntaxOk: boolean, hasFlowSpec: boolean }} inspectResult + */ +function flowInspectAllowsRun(inspectResult) { + return !!(inspectResult && inspectResult.syntaxOk && inspectResult.hasFlowSpec); +} + +/** + * Validates the active editor: must exist and be a Python document. + * Shows an error message and returns null when validation fails. + * @returns {import('vscode').TextDocument | null} + */ +function validateEditorForFlowCommands(editor) { + const vscode = require('vscode'); + + if (!editor) { + vscode.window.showErrorMessage('No active editor.'); + return null; + } + + const doc = editor.document; + if (doc.languageId !== 'python') { + vscode.window.showErrorMessage('Active file is not a Python file.'); + return null; + } + + return doc; +} + +/** + * Shows errors when inspect result blocks run/spin. + * @returns {boolean} true if run/spin should abort (failure reported), false if OK to continue. + * @param {{ syntaxOk: boolean, hasFlowSpec: boolean, error?: string | null }} inspectResult + */ +function reportFlowInspectFailure(inspectResult) { + const vscode = require('vscode'); + + if (flowInspectAllowsRun(inspectResult)) { + return false; + } + + if (!inspectResult.syntaxOk) { + vscode.window.showErrorMessage( + inspectResult.error || 'This file has a Python syntax error.' + ); + return true; + } + + if (!inspectResult.hasFlowSpec) { + vscode.window.showErrorMessage( + 'No FlowSpec class found. This file does not appear to be a Metaflow flow.' + ); + return true; + } + + return true; +} + +module.exports = { + flowInspectAllowsRun, + validateEditorForFlowCommands, + reportFlowInspectFailure, +}; diff --git a/package.json b/package.json index 31a7c2f..dbe89e9 100644 --- a/package.json +++ b/package.json @@ -13,6 +13,9 @@ "url": "https://github.com/outerbounds/metaflow-dev-vscode" }, "main": "./extension.js", + "scripts": { + "test": "node --test test/*.test.js" + }, "contributes": { "commands": [ { diff --git a/parameterPrompt.js b/parameterPrompt.js new file mode 100644 index 0000000..3e711ea --- /dev/null +++ b/parameterPrompt.js @@ -0,0 +1,122 @@ +const vscode = require('vscode'); +const path = require('path'); +const { execFilePythonChain } = require('./pythonRunner'); + +const INSPECT_SCRIPT = path.join(__dirname, 'scripts', 'inspectFlowFile.py'); + +function parseInspectStdout(stdout) { + try { + const parsed = JSON.parse(stdout); + return { + syntaxOk: !!parsed.syntaxOk, + hasFlowSpec: !!parsed.hasFlowSpec, + parameters: Array.isArray(parsed.parameters) ? parsed.parameters : [], + error: parsed.error != null ? parsed.error : null, + }; + } catch (parseErr) { + console.error('Failed to parse inspect output:', parseErr.message); + return { + syntaxOk: false, + hasFlowSpec: false, + parameters: [], + error: 'Invalid inspect output.', + }; + } +} + +/** + * Run inspectFlowFile.py and return { syntaxOk, hasFlowSpec, parameters, error }. + * Tries `python` then `python3` when PYTHON is unset. On failure, returns a result object (never throws). + */ +function inspectFlowFile(filePath) { + return new Promise((resolve) => { + execFilePythonChain( + [INSPECT_SCRIPT, filePath], + { timeout: 10000 }, + (err, stdout, stderr) => { + if (err) { + console.error('inspectFlowFile failed:', stderr || err.message); + resolve({ + syntaxOk: false, + hasFlowSpec: false, + parameters: [], + error: (stderr && String(stderr).trim()) || err.message || 'Failed to inspect flow file.', + }); + return; + } + resolve(parseInspectStdout(stdout)); + } + ); + }); +} + +/** + * Spawn the Python AST extractor and return parsed parameter definitions. + * Returns an empty array on failure (logged). + */ +function extractParameters(filePath) { + return inspectFlowFile(filePath).then((inspect) => inspect.parameters); +} + +/** + * Validate user input against the declared parameter type. + */ +function validateInput(value, paramType, required) { + if (!value.trim()) { + return required ? 'This parameter is required.' : undefined; + } + const trimmed = value.trim(); + switch (paramType) { + case 'int': + return /^-?\d+$/.test(trimmed) ? undefined : 'Expected an integer.'; + case 'float': + return isNaN(Number(trimmed)) ? 'Expected a number.' : undefined; + case 'bool': + return /^(true|false|0|1)$/i.test(trimmed) ? undefined : 'Expected true or false.'; + case 'JSONType': + try { JSON.parse(trimmed); return undefined; } catch { return 'Expected valid JSON.'; } + default: + return undefined; + } +} + +/** + * Prompt the user for values of each extracted parameter. + * Returns an array of {name, value} pairs, or null if the user cancelled. + */ +async function promptForParameters(params) { + const results = []; + + for (const param of params) { + const defaultStr = param.default !== null && param.default !== undefined + ? String(param.default) + : ''; + + const typeHint = param.type ? ` (${param.type})` : ''; + const requiredHint = param.required ? ' [required]' : ''; + const helpText = param.help ? ` — ${param.help}` : ''; + + const value = await vscode.window.showInputBox({ + title: `Metaflow Parameter: --${param.name}`, + prompt: `${param.name}${typeHint}${requiredHint}${helpText}`, + value: defaultStr, + validateInput: (input) => validateInput(input, param.type, param.required), + }); + + if (value === undefined) { + return null; + } + + // Preserve user input (including empty string and intentional spaces) for CLI flags. + results.push({ name: param.name, value }); + } + + return results; +} + +module.exports = { + inspectFlowFile, + extractParameters, + promptForParameters, + validateInput, +}; diff --git a/pythonRunner.js b/pythonRunner.js new file mode 100644 index 0000000..f526cde --- /dev/null +++ b/pythonRunner.js @@ -0,0 +1,77 @@ +const { execFile, execFileSync } = require('child_process'); + +/** + * Interpreters to try when PYTHON is unset: many Linux images only ship `python3`. + */ +const DEFAULT_PYTHON_CANDIDATES = ['python', 'python3']; + +function isMissingInterpreterError(err) { + return err && (err.code === 'ENOENT' || err.errno === -4058); +} + +/** + * Ordered list of Python executables: env PYTHON, then defaults. + */ +function pythonCandidates() { + if (process.env.PYTHON && String(process.env.PYTHON).trim()) { + return [process.env.PYTHON.trim()]; + } + return DEFAULT_PYTHON_CANDIDATES.slice(); +} + +/** + * Run a Python binary with args; on missing interpreter, try next candidate. + * @param {string[]} args + * @param {import('child_process').ExecFileOptionsWithStringEncoding} opts + * @param {(err: Error | null, stdout: string, stderr: string) => void} callback + */ +function execFilePythonChain(args, opts, callback) { + const chain = pythonCandidates(); + function attempt(i) { + if (i >= chain.length) { + callback( + new Error( + 'No Python interpreter found. Install Python or set the PYTHON environment variable.' + ), + '', + '' + ); + return; + } + execFile(chain[i], args, opts, (err, stdout, stderr) => { + if (err && isMissingInterpreterError(err)) { + attempt(i + 1); + return; + } + callback(err, stdout, stderr); + }); + } + attempt(0); +} + +/** + * Sync variant for tests: PYTHON, then python, then python3. + * @param {string[]} args + * @param {import('child_process').ExecSyncOptionsWithStringEncoding} opts + */ +function execFilePythonSync(args, opts) { + const chain = pythonCandidates(); + let lastErr; + for (let i = 0; i < chain.length; i++) { + try { + return execFileSync(chain[i], args, opts); + } catch (err) { + lastErr = err; + if (isMissingInterpreterError(err) && i < chain.length - 1) continue; + throw err; + } + } + throw lastErr; +} + +module.exports = { + pythonCandidates, + execFilePythonChain, + execFilePythonSync, + isMissingInterpreterError, +}; diff --git a/scripts/extractFlowParameters.py b/scripts/extractFlowParameters.py new file mode 100644 index 0000000..dbd6b73 --- /dev/null +++ b/scripts/extractFlowParameters.py @@ -0,0 +1,44 @@ +""" +Extract Metaflow Parameter definitions from a Python flow file using AST only. +Never imports or executes user code. + +Usage: python extractFlowParameters.py + +Outputs a JSON array to stdout, e.g.: +[ + {"name": "epochs", "attribute": "epochs", "default": 10, "type": "int", "required": false, "help": "Number of training epochs"}, + {"name": "lr", "attribute": "learning_rate", "default": null, "type": "float", "required": true, "help": null} +] +""" + +import json +import sys + +from metaflow_flow_ast import extract_parameters + + +def main(): + if len(sys.argv) != 2: + print('Usage: python extractFlowParameters.py ', file=sys.stderr) + sys.exit(1) + + filepath = sys.argv[1] + try: + with open(filepath, 'r', encoding='utf-8') as f: + source = f.read() + except (OSError, IOError) as e: + print(f'Cannot read file: {e}', file=sys.stderr) + sys.exit(1) + + try: + params = extract_parameters(source) + except SyntaxError as e: + print(f'Syntax error in {filepath}: {e}', file=sys.stderr) + sys.exit(1) + + json.dump(params, sys.stdout, indent=2) + sys.stdout.write('\n') + + +if __name__ == '__main__': + main() diff --git a/scripts/inspectFlowFile.py b/scripts/inspectFlowFile.py new file mode 100644 index 0000000..1874f44 --- /dev/null +++ b/scripts/inspectFlowFile.py @@ -0,0 +1,55 @@ +""" +Emit a single JSON object describing syntax, FlowSpec presence, and parameters. + +Usage: python inspectFlowFile.py + +Always exits 0 when JSON is printed to stdout so callers can parse reliably. +""" + +import json +import sys + +from metaflow_flow_ast import analyze_flow_source + + +def main(): + if len(sys.argv) != 2: + result = { + 'syntaxOk': False, + 'hasFlowSpec': False, + 'parameters': [], + 'error': 'Usage: python inspectFlowFile.py ', + } + json.dump(result, sys.stdout, indent=2) + sys.stdout.write('\n') + sys.exit(0) + + filepath = sys.argv[1] + try: + with open(filepath, 'r', encoding='utf-8') as f: + source = f.read() + except (OSError, IOError) as e: + result = { + 'syntaxOk': False, + 'hasFlowSpec': False, + 'parameters': [], + 'error': f'Cannot read file: {e}', + } + json.dump(result, sys.stdout, indent=2) + sys.stdout.write('\n') + sys.exit(0) + + analyzed = analyze_flow_source(source) + # Normalize: ensure error key exists when syntaxOk is False + out = { + 'syntaxOk': analyzed['syntaxOk'], + 'hasFlowSpec': analyzed['hasFlowSpec'], + 'parameters': analyzed['parameters'], + 'error': analyzed.get('error'), + } + json.dump(out, sys.stdout, indent=2) + sys.stdout.write('\n') + + +if __name__ == '__main__': + main() diff --git a/scripts/metaflow_flow_ast.py b/scripts/metaflow_flow_ast.py new file mode 100644 index 0000000..9418db6 --- /dev/null +++ b/scripts/metaflow_flow_ast.py @@ -0,0 +1,168 @@ +""" +AST-only helpers for Metaflow flow files. Never imports or executes user code. +""" + +from __future__ import annotations + +import ast +from typing import Any, Optional, Tuple + + +def _get_call_name(node: ast.AST) -> Optional[str]: + """Return the simple name of a Call target, e.g. 'Parameter' or None.""" + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return None + + +def _get_literal(node: ast.AST) -> Any: + """Return a JSON-safe literal from an AST Constant node, or None.""" + if isinstance(node, ast.Constant) and isinstance(node.value, (str, int, float, bool)): + return node.value + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + inner = _get_literal(node.operand) + if isinstance(inner, (int, float)): + return -inner + return None + + +def _get_type_name(node: ast.AST) -> Optional[str]: + """Return the type name string from a simple Name node like str/int/float/bool.""" + if isinstance(node, ast.Name) and node.id in ('str', 'int', 'float', 'bool', 'JSONType'): + return node.id + return None + + +def _is_flowspec_base(base: ast.AST) -> bool: + """Check whether a base class node refers to FlowSpec.""" + if isinstance(base, ast.Name) and base.id == 'FlowSpec': + return True + if isinstance(base, ast.Attribute) and base.attr == 'FlowSpec': + return True + return False + + +def has_flowspec_subclass(tree: ast.AST) -> bool: + """True if the module defines at least one class inheriting from FlowSpec.""" + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and any(_is_flowspec_base(b) for b in node.bases): + return True + return False + + +def _extract_parameter(call_node: ast.Call, attribute_name: str) -> dict[str, Any]: + """Extract parameter metadata from a Parameter(...) call AST node.""" + param: dict[str, Any] = { + 'name': None, + 'attribute': attribute_name, + 'default': None, + 'type': None, + 'required': False, + 'help': None, + } + + if call_node.args: + name_val = _get_literal(call_node.args[0]) + if isinstance(name_val, str): + param['name'] = name_val + + if param['name'] is None: + param['name'] = attribute_name + + for kw in call_node.keywords: + if kw.arg == 'default': + param['default'] = _get_literal(kw.value) + elif kw.arg == 'type': + param['type'] = _get_type_name(kw.value) + elif kw.arg == 'required': + if isinstance(kw.value, ast.Constant) and kw.value.value is True: + param['required'] = True + elif kw.arg == 'help': + help_val = _get_literal(kw.value) + if isinstance(help_val, str): + param['help'] = help_val + + return param + + +def _parameter_from_class_body_item(item: ast.stmt) -> Optional[Tuple[ast.Call, str]]: + """ + Match `x = Parameter(...)` or annotated `x: T = Parameter(...)`. + Returns (call_node, attribute_name) or None. + """ + if isinstance(item, ast.Assign): + if len(item.targets) != 1 or not isinstance(item.targets[0], ast.Name): + return None + if not isinstance(item.value, ast.Call): + return None + if _get_call_name(item.value.func) != 'Parameter': + return None + return (item.value, item.targets[0].id) + if isinstance(item, ast.AnnAssign): + if not isinstance(item.target, ast.Name): + return None + if not isinstance(item.value, ast.Call): + return None + if _get_call_name(item.value.func) != 'Parameter': + return None + return (item.value, item.target.id) + return None + + +def extract_parameters_from_tree(tree: ast.AST) -> list[dict[str, Any]]: + """Collect Parameter(...) definitions from FlowSpec subclasses in a parsed tree.""" + parameters: list[dict[str, Any]] = [] + + for node in ast.walk(tree): + if not isinstance(node, ast.ClassDef): + continue + if not any(_is_flowspec_base(b) for b in node.bases): + continue + + for item in node.body: + got = _parameter_from_class_body_item(item) + if got is None: + continue + call_node, attr_name = got + parameters.append(_extract_parameter(call_node, attr_name)) + + return parameters + + +def extract_parameters(source: str) -> list[dict[str, Any]]: + """Parse source and return Parameter definitions from FlowSpec classes.""" + tree = ast.parse(source) + return extract_parameters_from_tree(tree) + + +def analyze_flow_source(source: str) -> dict[str, Any]: + """ + Single parse: syntax validity, FlowSpec presence, and parameters. + + Returns keys: syntaxOk (bool), hasFlowSpec (bool), parameters (list), + and error (str | None) when syntaxOk is False. + """ + try: + tree = ast.parse(source) + except SyntaxError as e: + msg = e.msg or 'invalid syntax' + line = e.lineno + err = f'Syntax error: {msg}' + (f' (line {line})' if line else '') + return { + 'syntaxOk': False, + 'hasFlowSpec': False, + 'parameters': [], + 'error': err, + } + + has_flow = has_flowspec_subclass(tree) + params = extract_parameters_from_tree(tree) + + return { + 'syntaxOk': True, + 'hasFlowSpec': has_flow, + 'parameters': params, + 'error': None, + } diff --git a/test/buildRunCommand.test.js b/test/buildRunCommand.test.js new file mode 100644 index 0000000..5998f19 --- /dev/null +++ b/test/buildRunCommand.test.js @@ -0,0 +1,152 @@ +const { describe, it } = require('node:test'); +const assert = require('node:assert/strict'); +const { + buildRunCommand, quoteForTerminal, quoteForBash, quoteForPowershell, + quoteForCmd, detectShellType, CLI_NAME_RE, +} = require('../buildRunCommand'); + +describe('quoteForBash', () => { + it('wraps in single quotes', () => { + assert.equal(quoteForBash('hello'), "'hello'"); + }); + + it('passes through $, `, \\ literally inside single quotes', () => { + assert.equal(quoteForBash('$HOME'), "'$HOME'"); + assert.equal(quoteForBash('`cmd`'), "'`cmd`'"); + assert.equal(quoteForBash('C:\\Users'), "'C:\\Users'"); + }); + + it('escapes single quotes with close-escape-reopen pattern', () => { + assert.equal(quoteForBash("it's"), "'it'\\''s'"); + }); +}); + +describe('quoteForPowershell', () => { + it('wraps in single quotes', () => { + assert.equal(quoteForPowershell('hello'), "'hello'"); + }); + + it('passes through $, `, \\ literally inside single quotes', () => { + assert.equal(quoteForPowershell('$HOME'), "'$HOME'"); + assert.equal(quoteForPowershell('`cmd`'), "'`cmd`'"); + assert.equal(quoteForPowershell('C:\\Users\\me'), "'C:\\Users\\me'"); + }); + + it('escapes single quotes by doubling', () => { + assert.equal(quoteForPowershell("it's"), "'it''s'"); + }); +}); + +describe('quoteForCmd', () => { + it('wraps in double quotes', () => { + assert.equal(quoteForCmd('hello'), '"hello"'); + }); + + it('escapes double quotes by doubling', () => { + assert.equal(quoteForCmd('say "hi"'), '"say ""hi"""'); + }); + + it('escapes percent signs by doubling', () => { + assert.equal(quoteForCmd('%PATH%'), '"%%PATH%%"'); + }); +}); + +describe('quoteForTerminal dispatches by shell', () => { + it('uses bash quoting by default', () => { + assert.equal(quoteForTerminal('$x', 'bash'), "'$x'"); + }); + + it('uses powershell quoting when specified', () => { + assert.equal(quoteForTerminal("it's", 'powershell'), "'it''s'"); + }); + + it('uses cmd quoting when specified', () => { + assert.equal(quoteForTerminal('%PATH%', 'cmd'), '"%%PATH%%"'); + }); +}); + +describe('detectShellType', () => { + it('detects PowerShell', () => { + assert.equal(detectShellType('C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe'), 'powershell'); + assert.equal(detectShellType('/usr/local/bin/pwsh'), 'powershell'); + }); + + it('detects cmd.exe', () => { + assert.equal(detectShellType('C:\\Windows\\System32\\cmd.exe'), 'cmd'); + }); + + it('defaults to bash for other shells', () => { + assert.equal(detectShellType('/bin/bash'), 'bash'); + assert.equal(detectShellType('/bin/zsh'), 'bash'); + assert.equal(detectShellType(''), 'bash'); + assert.equal(detectShellType(undefined), 'bash'); + }); +}); + +describe('CLI_NAME_RE', () => { + it('accepts valid Python identifiers', () => { + assert.ok(CLI_NAME_RE.test('epochs')); + assert.ok(CLI_NAME_RE.test('learning_rate')); + assert.ok(CLI_NAME_RE.test('_private')); + }); + + it('rejects invalid names', () => { + assert.ok(!CLI_NAME_RE.test('')); + assert.ok(!CLI_NAME_RE.test('123abc')); + assert.ok(!CLI_NAME_RE.test('my param')); + assert.ok(!CLI_NAME_RE.test('a=b')); + }); + + it('accepts hyphenated Click-style names', () => { + assert.ok(CLI_NAME_RE.test('my-param')); + assert.ok(CLI_NAME_RE.test('foo-bar_baz')); + }); +}); + +describe('buildRunCommand', () => { + it('builds a bare run command with no flags (bash)', () => { + const cmd = buildRunCommand('/path/to/flow.py', [], 'bash'); + assert.equal(cmd, "python '/path/to/flow.py' run"); + }); + + it('appends --name=value flags (bash)', () => { + const cmd = buildRunCommand('/flow.py', [ + { name: 'epochs', value: '20' }, + { name: 'lr', value: '0.001' }, + ], 'bash'); + assert.equal(cmd, "python '/flow.py' run --epochs='20' --lr='0.001'"); + }); + + it('builds command for powershell', () => { + const cmd = buildRunCommand('C:\\flow.py', [ + { name: 'model', value: "it's a model" }, + ], 'powershell'); + assert.equal(cmd, "python 'C:\\flow.py' run --model='it''s a model'"); + }); + + it('quotes values with spaces (bash)', () => { + const cmd = buildRunCommand('/flow.py', [ + { name: 'model', value: 'my model' }, + ], 'bash'); + assert.equal(cmd, "python '/flow.py' run --model='my model'"); + }); + + it('accepts hyphenated parameter names in command', () => { + const cmd = buildRunCommand('/flow.py', [{ name: 'my-flag', value: '1' }], 'bash'); + assert.equal(cmd, "python '/flow.py' run --my-flag='1'"); + }); + + it('throws on invalid parameter names', () => { + assert.throws( + () => buildRunCommand('/flow.py', [{ name: 'bad=name', value: '1' }], 'bash'), + /Invalid parameter name/ + ); + }); + + it('throws on newlines in values', () => { + assert.throws( + () => buildRunCommand('/flow.py', [{ name: 'x', value: 'a\nb' }], 'bash'), + /newline/ + ); + }); +}); diff --git a/test/extractParameters.test.js b/test/extractParameters.test.js new file mode 100644 index 0000000..95b5773 --- /dev/null +++ b/test/extractParameters.test.js @@ -0,0 +1,61 @@ +const { describe, it } = require('node:test'); +const assert = require('node:assert/strict'); +const path = require('path'); +const { execFilePythonSync } = require('../pythonRunner'); + +const SCRIPT = path.join(__dirname, '..', 'scripts', 'extractFlowParameters.py'); +const FIXTURES = path.join(__dirname, '..', 'test_fixtures'); + +function runExtractor(fixture) { + const stdout = execFilePythonSync([SCRIPT, path.join(FIXTURES, fixture)], { + encoding: 'utf-8', + }); + return JSON.parse(stdout); +} + +describe('extractFlowParameters.py', () => { + it('extracts parameters from a typical flow', () => { + const params = runExtractor('sample_flow.py'); + assert.equal(params.length, 4); + + const epochs = params.find((p) => p.name === 'epochs'); + assert.equal(epochs.attribute, 'epochs'); + assert.equal(epochs.default, 10); + assert.equal(epochs.type, 'int'); + assert.equal(epochs.help, 'Number of training epochs'); + assert.equal(epochs.required, false); + + const lr = params.find((p) => p.name === 'lr'); + assert.equal(lr.attribute, 'learning_rate'); + assert.equal(lr.default, 0.001); + assert.equal(lr.type, 'float'); + + const model = params.find((p) => p.name === 'model'); + assert.equal(model.required, true); + assert.equal(model.default, null); + + const verbose = params.find((p) => p.name === 'verbose'); + assert.equal(verbose.default, true); + assert.equal(verbose.type, 'bool'); + }); + + it('returns empty array for non-flow files', () => { + const params = runExtractor('not_a_flow.py'); + assert.deepEqual(params, []); + }); + + it('returns empty array for flow without parameters', () => { + const params = runExtractor('no_params_flow.py'); + assert.deepEqual(params, []); + }); + + it('extracts Parameter from annotated assignment', () => { + const params = runExtractor('ann_assign_flow.py'); + assert.equal(params.length, 1); + const epochs = params[0]; + assert.equal(epochs.name, 'epochs'); + assert.equal(epochs.attribute, 'epochs'); + assert.equal(epochs.default, 5); + assert.equal(epochs.type, 'int'); + }); +}); diff --git a/test/flowDetection.test.js b/test/flowDetection.test.js new file mode 100644 index 0000000..f0358c4 --- /dev/null +++ b/test/flowDetection.test.js @@ -0,0 +1,22 @@ +const { describe, it } = require('node:test'); +const assert = require('node:assert/strict'); +const { flowInspectAllowsRun } = require('../flowDetection'); + +describe('flowInspectAllowsRun', () => { + it('allows when syntaxOk and hasFlowSpec', () => { + assert.ok(flowInspectAllowsRun({ syntaxOk: true, hasFlowSpec: true })); + }); + + it('blocks when syntaxOk is false', () => { + assert.ok(!flowInspectAllowsRun({ syntaxOk: false, hasFlowSpec: false })); + }); + + it('blocks when hasFlowSpec is false', () => { + assert.ok(!flowInspectAllowsRun({ syntaxOk: true, hasFlowSpec: false })); + }); + + it('blocks on null or missing fields', () => { + assert.ok(!flowInspectAllowsRun(null)); + assert.ok(!flowInspectAllowsRun({})); + }); +}); diff --git a/test/inspectFlowFile.test.js b/test/inspectFlowFile.test.js new file mode 100644 index 0000000..6e598ad --- /dev/null +++ b/test/inspectFlowFile.test.js @@ -0,0 +1,48 @@ +const { describe, it } = require('node:test'); +const assert = require('node:assert/strict'); +const path = require('path'); +const { execFilePythonSync } = require('../pythonRunner'); + +const SCRIPT = path.join(__dirname, '..', 'scripts', 'inspectFlowFile.py'); +const FIXTURES = path.join(__dirname, '..', 'test_fixtures'); + +function inspectFixture(name) { + const stdout = execFilePythonSync([SCRIPT, path.join(FIXTURES, name)], { + encoding: 'utf-8', + }); + return JSON.parse(stdout); +} + +describe('inspectFlowFile.py', () => { + it('returns syntaxOk and hasFlowSpec for sample flow with parameters', () => { + const j = inspectFixture('sample_flow.py'); + assert.equal(j.syntaxOk, true); + assert.equal(j.hasFlowSpec, true); + assert.equal(j.parameters.length, 4); + assert.equal(j.error, null); + }); + + it('returns hasFlowSpec false for non-flow file', () => { + const j = inspectFixture('not_a_flow.py'); + assert.equal(j.syntaxOk, true); + assert.equal(j.hasFlowSpec, false); + assert.deepEqual(j.parameters, []); + }); + + it('returns syntaxOk false for invalid Python', () => { + const j = inspectFixture('syntax_error_flow.py'); + assert.equal(j.syntaxOk, false); + assert.equal(j.hasFlowSpec, false); + assert.ok(j.error && j.error.includes('Syntax error')); + }); + + it('matches extractParameters list for sample_flow', () => { + const extractScript = path.join(__dirname, '..', 'scripts', 'extractFlowParameters.py'); + const extractOut = execFilePythonSync([extractScript, path.join(FIXTURES, 'sample_flow.py')], { + encoding: 'utf-8', + }); + const paramsExtract = JSON.parse(extractOut); + const j = inspectFixture('sample_flow.py'); + assert.deepEqual(j.parameters, paramsExtract); + }); +}); diff --git a/test_fixtures/ann_assign_flow.py b/test_fixtures/ann_assign_flow.py new file mode 100644 index 0000000..8fb3663 --- /dev/null +++ b/test_fixtures/ann_assign_flow.py @@ -0,0 +1,18 @@ +from metaflow import FlowSpec, Parameter, step + + +class AnnFlow(FlowSpec): + epochs: int = Parameter('epochs', default=5, type=int) + + @step + def start(self): + print(self.epochs) + self.next(self.end) + + @step + def end(self): + pass + + +if __name__ == '__main__': + AnnFlow() diff --git a/test_fixtures/no_params_flow.py b/test_fixtures/no_params_flow.py new file mode 100644 index 0000000..15f03b8 --- /dev/null +++ b/test_fixtures/no_params_flow.py @@ -0,0 +1,15 @@ +from metaflow import FlowSpec, step + + +class SimpleFlow(FlowSpec): + @step + def start(self): + self.next(self.end) + + @step + def end(self): + pass + + +if __name__ == '__main__': + SimpleFlow() diff --git a/test_fixtures/not_a_flow.py b/test_fixtures/not_a_flow.py new file mode 100644 index 0000000..bb92f10 --- /dev/null +++ b/test_fixtures/not_a_flow.py @@ -0,0 +1,5 @@ +class SomeClass: + x = 10 + +def hello(): + print("not a flow") diff --git a/test_fixtures/sample_flow.py b/test_fixtures/sample_flow.py new file mode 100644 index 0000000..785c941 --- /dev/null +++ b/test_fixtures/sample_flow.py @@ -0,0 +1,21 @@ +from metaflow import FlowSpec, Parameter, step + + +class TrainFlow(FlowSpec): + epochs = Parameter('epochs', default=10, type=int, help='Number of training epochs') + learning_rate = Parameter('lr', default=0.001, type=float) + model_name = Parameter('model', required=True, help='Name of the model to train') + verbose = Parameter('verbose', default=True, type=bool) + + @step + def start(self): + print(f"Training {self.model_name} for {self.epochs} epochs") + self.next(self.end) + + @step + def end(self): + pass + + +if __name__ == '__main__': + TrainFlow() diff --git a/test_fixtures/syntax_error_flow.py b/test_fixtures/syntax_error_flow.py new file mode 100644 index 0000000..22a303b --- /dev/null +++ b/test_fixtures/syntax_error_flow.py @@ -0,0 +1,11 @@ +from metaflow import FlowSpec, step + + +class BrokenFlow(FlowSpec): + @step + def start(self): + self.next(self.end + + @step + def end(self): + pass