Skip to content

Commit

Permalink
Generate module from python cells in mojo context
Browse files Browse the repository at this point in the history
  • Loading branch information
@jackos
jackos committed Jul 3, 2024
1 parent cae5926 commit fa57030
Show file tree
Hide file tree
Showing 5 changed files with 66 additions and 52 deletions.
13 changes: 0 additions & 13 deletions .vscode/settings.json
View file

This file was deleted.

3 changes: 3 additions & 0 deletions CHANGELOG.md
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
# Changelog

## v0.6.7
- Auto generate module from python cells when running Mojo cell

## v0.6.6
- Default to using MDL.AI when opening a .md file
- Improve README.md
Expand Down
31 changes: 18 additions & 13 deletions src/kernel.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,8 @@ export class Kernel {
// Build a object containing languages and their cells
let cellsStripped: Cell[] = [];
let matchingCells = 0;
let pythonCells = 0;
let pythonMatchingCells = 0;
let pythonCells: Cell[] = [];
for (const cell of cellsUpToCurrent) {
if (cell.document.languageId === cells[0].document.languageId) {
matchingCells++;
Expand All @@ -64,9 +65,15 @@ export class Kernel {
cell: cell,
});
}
// Also capture python cells if they exist when running Mojo
if (cells[0].document.languageId === "mojo") {
if(cell.document.languageId === "python") {
pythonCells += 1;
if (cell.document.languageId === "python") {
pythonMatchingCells++
pythonCells.push({
index: pythonMatchingCells,
contents: cell.document.getText(),
cell: cell,
});
}
}
}
Expand All @@ -86,7 +93,7 @@ export class Kernel {
'Authorization': `Bearer ${getGroqAIKey()}`,
}

const messages: ChatMessage[] = [{"role": "user", "content": "You're generating codeblocks to help users solve programming problems, make sure that you put the name of the language in the markdown blocks like ```python"}]
const messages: ChatMessage[] = [{ "role": "user", "content": "You're generating codeblocks to help users solve programming problems, make sure that you put the name of the language in the markdown blocks like ```python" }]
for (const message of cellsStripped) {
messages.push({ role: "user", content: message.contents });
}
Expand Down Expand Up @@ -210,7 +217,7 @@ export class Kernel {
return
}
lastRunLanguage = "mojo";
let mojoResult = processCellsMojo(cellsStripped);
let mojoResult = processCellsMojo(cellsStripped, pythonCells);
output = mojoResult.stream
clearOutput = mojoResult.clearOutput
break;
Expand All @@ -231,14 +238,13 @@ export class Kernel {
output = processCellsGo(cellsStripped);
break;
case "python":
let command = "python"
let command = "python3"
if (commandNotOnPath(command, "")) {
if (commandNotOnPath("python3", "https://www.python.org/downloads/")) {
if (commandNotOnPath("python", "https://www.python.org/downloads/")) {
exec.end(false, (new Date).getTime());
return
} else {
command = "python3"
}
command = "python"
}
lastRunLanguage = "python";
let pyResult = processCellsPython(cellsStripped, command);
Expand Down Expand Up @@ -343,7 +349,7 @@ export class Kernel {
if (lastRunLanguage == "shell") {
currentCellOutput = outputs[1]
} else {
currentCellOutput = outputs[currentCellLang.index + pythonCells];
currentCellOutput = outputs[currentCellLang.index + pythonCells.length];
}
if (!clearOutput && currentCellOutput.trim()) {
exec.replaceOutput([new NotebookCellOutput([NotebookCellOutputItem.text(currentCellOutput)])]);
Expand All @@ -360,7 +366,6 @@ export class Kernel {
}

// Loop through all the cells and increment version of image if it exists

if (doc.getCells().length >= (cells[0].index + 1)) {
let cell = doc.getCells(new NotebookRange(cells[0].index + 1, cells[0].index + 2))[0]
if (cell.kind === vscode.NotebookCellKind.Markup) {
Expand Down Expand Up @@ -392,8 +397,8 @@ export class Kernel {
}).then(() => {
// Execute commands to toggle cell edit mode and then toggle it back to preview.
vscode.commands.executeCommand('notebook.cell.edit').then(() => {
vscode.commands.executeCommand('notebook.cell.quitEdit').then(() => {
// Optionally, add any additional logic that needs to run after the refresh.
vscode.commands.executeCommand('notebook.cell.quitEdit').then(() => {
// Optionally, add any additional logic that needs to run after the refresh.
});
});
});
Expand Down
45 changes: 32 additions & 13 deletions src/languages/mojo.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,27 +4,37 @@ import { getTempPath, modularHome } from "../config";
import { Cell, CommentDecorator } from "../types";
import path from "path";
import { window } from "vscode";
import { processCellsPython } from "./python";
import { commandNotOnPath } from "../utils";

let tempDir = getTempPath();

export let processCellsMojo = (cells: Cell[]): { stream: ChildProcessWithoutNullStreams, clearOutput: boolean } => {
export let processCellsMojo = (cells: Cell[], pythonCells: Cell[]): { stream: ChildProcessWithoutNullStreams, clearOutput: boolean } => {
// If any python cells exist, make sure the generated file is current
if (pythonCells) {
let command = "python3"
if (commandNotOnPath(command, "")) {
command = "python"
}
processCellsPython(pythonCells, command)
}
const activeFilePath = path.dirname(window.activeTextEditor?.document.uri.fsPath as string);
let outerScope = "";

let innerScope = `def main():`;

let pythonFileExists = existsSync(path.join(tempDir, "md_notebook.py"));
let pythonFileExists = existsSync(path.join(tempDir, "mdl.py"));
if (pythonFileExists) {
outerScope += "from python import Python\n"
innerScope += `\n sys = Python.import_module("sys")\n sys.path.append("${activeFilePath}")\n sys.path.append("${tempDir}")\n py = Python.import_module("md_notebook")\n`
innerScope += `\n sys = Python.import_module("sys")\n sys.path.append("${activeFilePath}")\n sys.path.append("${tempDir}")\n py = Python.import_module("mdl")\n`
}
let cellCount = 0;
let clearOutput = false;
let inOuterScope = false;

for (const cell of cells) {
innerScope += `\n print("!!output-start-cell")\n`;
cell.contents = cell.contents.trim();
// cell.contents = cell.contents.trim();
const regex = /(\s*print\s*\()(.*?)(\)\s*$)/gm;

cell.contents = cell.contents.replace(regex, (_, before, content, after) => {
Expand All @@ -40,15 +50,24 @@ export let processCellsMojo = (cells: Cell[]): { stream: ChildProcessWithoutNull
});

cellCount++;
if(cell.contents.startsWith("#mdl:skip") || cell.contents.startsWith("# mdl:skip")) {
if (cell.contents.startsWith("#mdl:skip") || cell.contents.startsWith("# mdl:skip")) {
continue;
}
}
let lines = cell.contents.split("\n");
const len = lines.length;
let i = 0
for (let line of lines) {
i++
if(line.startsWith("struct") || line.startsWith("trait") || line.startsWith("from")) {
// Keep things in outerScope if they should not go in main()
if (
line.startsWith("struct") ||
line.startsWith("trait") ||
line.startsWith("fn") ||
line.startsWith("alias") ||
line.startsWith("from") ||
line.startsWith("import") ||
line.startsWith("@")
) {
inOuterScope = true;
} else if (!line.startsWith(" ") && !(line === "")) {
inOuterScope = false;
Expand Down Expand Up @@ -76,7 +95,7 @@ export let processCellsMojo = (cells: Cell[]): { stream: ChildProcessWithoutNull
if (line.startsWith("fn main():") || line.startsWith("def main():")) {
continue;
}
if (pythonFileExists && (line.includes('Python.import_module("sys")') || line.trim() == "from python import Python")){
if (pythonFileExists && (line.includes('Python.import_module("sys")') || line.trim() == "from python import Python")) {
continue;
}

Expand All @@ -90,7 +109,7 @@ export let processCellsMojo = (cells: Cell[]): { stream: ChildProcessWithoutNull
}

}
if(inOuterScope) {
if (inOuterScope) {
outerScope += line + "\n"
} else {
innerScope += " " + line + "\n";
Expand All @@ -105,9 +124,9 @@ export let processCellsMojo = (cells: Cell[]): { stream: ChildProcessWithoutNull
mkdirSync(tempDir, { recursive: true });
writeFileSync(mainFile, outerScope + innerScope);
let env = process.env;
if(typeof modularHome === "string") {
env = {"MODULAR_HOME": modularHome, ...env}
if (typeof modularHome === "string") {
env = { "MODULAR_HOME": modularHome, ...env }
}
return { stream: spawn('mojo', [mainFile], {cwd: activeFilePath, env}), clearOutput };

return { stream: spawn('mojo', [mainFile], { cwd: activeFilePath, env }), clearOutput };
};
26 changes: 13 additions & 13 deletions src/languages/python.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,12 @@ import { ChildProcessWithoutNullStreams, spawn } from "child_process";
import { mkdirSync, writeFileSync } from "fs";
import { getTempPath } from "../config";
import { Cell, CommentDecorator } from "../types";
import vscode from "vscode"
import vscode from "vscode"
import path from "path"

let tempDir = getTempPath();

export let processCellsPython = (cells: Cell[], command: string): {stream: ChildProcessWithoutNullStreams, clearOutput: boolean }=> {
export let processCellsPython = (cells: Cell[], command: string): { stream: ChildProcessWithoutNullStreams, clearOutput: boolean } => {

let innerScope = "";
let cellCount = 0;
Expand All @@ -31,23 +31,23 @@ export let processCellsPython = (cells: Cell[], command: string): {stream: Child
return `${before}${content}${after}`;
});
cellCount++;
if(cell.contents.startsWith("#mdl:skip") || cell.contents.startsWith("# mdl:skip")) {
if (cell.contents.startsWith("#mdl:skip") || cell.contents.startsWith("# mdl:skip")) {
continue;
}
if(cell.contents.startsWith("#mdl:skip")) {
}
if (cell.contents.startsWith("#mdl:skip")) {
continue
}
}
let lines = cell.contents.split("\n");
const len = lines.length;
let i = 0
for (let line of lines) {
i++
if (i==1 && line.replace(/\s/g, "").substring(0, 6) == "#file:") {
if (i == 1 && line.replace(/\s/g, "").substring(0, 6) == "#file:") {
let file = line.split(":")[1].trim()
if (file != "main.py"){
if (file != "main.py") {
let cleaned = ""
for(let line2 of lines){
if(line2.trim() != 'print("!!output-start-cell", flush=True)'){
for (let line2 of lines) {
if (line2.trim() != 'print("!!output-start-cell", flush=True)') {
cleaned += line2 + "\n"
}
}
Expand Down Expand Up @@ -78,11 +78,11 @@ export let processCellsPython = (cells: Cell[], command: string): {stream: Child
}
};

let mainFile = path.join(tempDir, "md_notebook.py");
let mainFile = path.join(tempDir, "mdl.py");
let header = `import sys\nsys.path.append("${activeFilePath}")\nsys.path.append("${tempDir}")\nfrom builtins import *\n`

mkdirSync(tempDir, { recursive: true });
writeFileSync(mainFile, header + innerScope);
return {stream: spawn(command, [mainFile], {cwd: activeFilePath}), clearOutput};

return { stream: spawn(command, [mainFile], { cwd: activeFilePath }), clearOutput };
};

0 comments on commit fa57030

Please sign in to comment.