Skip to content

Commit

Permalink
fix: prevent empty code execution outputs after remount
Browse files Browse the repository at this point in the history
  • Loading branch information
@jeremyphilemon
jeremyphilemon committed Jan 8, 2025
1 parent 50bd460 commit 1934fe3
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 86 deletions.
15 changes: 13 additions & 2 deletions components/console.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import {
} from 'react';
import { ConsoleOutput } from './block';
import { cn } from '@/lib/utils';
import { useBlockSelector } from '@/hooks/use-block';

interface ConsoleProps {
consoleOutputs: Array<ConsoleOutput>;
Expand All @@ -21,6 +22,8 @@ export function Console({ consoleOutputs, setConsoleOutputs }: ConsoleProps) {
const [isResizing, setIsResizing] = useState(false);
const consoleEndRef = useRef<HTMLDivElement>(null);

const isBlockVisible = useBlockSelector((state) => state.isVisible);

const minHeight = 100;
const maxHeight = 800;

Expand Down Expand Up @@ -57,6 +60,12 @@ export function Console({ consoleOutputs, setConsoleOutputs }: ConsoleProps) {
consoleEndRef.current?.scrollIntoView({ behavior: 'smooth' });
}, [consoleOutputs]);

useEffect(() => {
if (!isBlockVisible) {
setConsoleOutputs([]);
}
}, [isBlockVisible, setConsoleOutputs]);

return consoleOutputs.length > 0 ? (
<>
<div
Expand Down Expand Up @@ -115,14 +124,16 @@ export function Console({ consoleOutputs, setConsoleOutputs }: ConsoleProps) {
consoleOutput.status,
) ? (
<div className="flex flex-row gap-2">
<div className="animate-spin size-fit self-center">
<div className="animate-spin size-fit self-center mb-auto mt-0.5">
<LoaderIcon />
</div>
<div className="text-muted-foreground">
{consoleOutput.status === 'in_progress'
? 'Initializing...'
: consoleOutput.status === 'loading_packages'
? 'Loading Packages...'
? consoleOutput.contents.map((content) =>
content.type === 'text' ? content.value : null,
)
: null}
</div>
</div>
Expand Down
132 changes: 48 additions & 84 deletions components/run-code-button.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,49 +6,17 @@ import {
useCallback,
useState,
useEffect,
memo,
} from 'react';
import type { ConsoleOutput, UIBlock } from './block';
import type { ConsoleOutput, ConsoleOutputContent, UIBlock } from './block';
import { Button } from './ui/button';
import { PlayIcon } from './icons';

function detectPythonImports(code: string, pyodide: any): Set<string> {
const imports = new Set<string>();

const importPatterns = [
/import\s+(\w+)(?:\s+as\s+\w+)?/g,
/from\s+(\w+(?:\.\w+)*)\s+import/g,
/import\s+(\w+(?:\.\w+)*)/g,
];

for (const pattern of importPatterns) {
let match: RegExpExecArray | null;

while (true) {
match = pattern.exec(code);
if (match === null) break;

const rootPackage = match[1].split('.')[0];
if (rootPackage) {
imports.add(rootPackage);
}
}
}

// Get standard libraries dynamically when initializing Pyodide
let standardLibs = new Set<string>();
if (pyodide) {
const stdLibModules = pyodide.runPython(`
import sys
list(sys.stdlib_module_names)
`);
standardLibs = new Set(stdLibModules);
}

return new Set(Array.from(imports).filter((pkg) => !standardLibs.has(pkg)));
}
import { useBlockSelector } from '@/hooks/use-block';

const OUTPUT_HANDLERS = {
matplotlib: `
import io
import base64
from matplotlib import pyplot as plt
# Clear any existing plots
Expand Down Expand Up @@ -91,19 +59,23 @@ function detectRequiredHandlers(code: string): string[] {
return handlers;
}

export function RunCodeButton({
block,
export function PureRunCodeButton({
setConsoleOutputs,
}: {
block: UIBlock;
setConsoleOutputs: Dispatch<SetStateAction<Array<ConsoleOutput>>>;
}) {
const [pyodide, setPyodide] = useState<any>(null);
const isPython = true;
const codeContent = block.content;
const [pyodide, setPyodide] = useState<any>(null);

const codeContent = useBlockSelector((state) => state.content);
const isBlockStreaming = useBlockSelector(
(state) => state.status === 'streaming',
);

const loadAndRunPython = useCallback(async () => {
const runId = generateUUID();
const stdOutputs: Array<ConsoleOutputContent> = [];

setConsoleOutputs((outputs) => [
...outputs,
Expand All @@ -120,41 +92,38 @@ export function RunCodeButton({
try {
if (!currentPyodideInstance) {
// @ts-expect-error - loadPyodide is not defined
const newPyodideInstance = await loadPyodide({
const newPyodideInstance = await globalThis.loadPyodide({
indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.23.4/full/',
});

setPyodide(null);
setPyodide(newPyodideInstance);
currentPyodideInstance = newPyodideInstance;
}

await currentPyodideInstance.runPythonAsync(`
import sys, io, gc
import base64
sys.stdout = io.StringIO()
`);

// Detect and load required packages
const requiredPackages = detectPythonImports(
codeContent,
currentPyodideInstance,
);

if (requiredPackages.size > 0) {
setConsoleOutputs((outputs) => [
...outputs.filter((output) => output.id !== runId),
{
id: runId,
contents: [],
status: 'loading_packages',
},
]);

await currentPyodideInstance.loadPackage(
Array.from(requiredPackages),
);
}
currentPyodideInstance.setStdout({
batched: (output: string) => {
stdOutputs.push({
type: output.startsWith('data:image/png;base64')
? 'image'
: 'text',
value: output,
});
},
});

await currentPyodideInstance.loadPackagesFromImports(codeContent, {
messageCallback: (message: string) => {
setConsoleOutputs((outputs) => [
...outputs.filter((output) => output.id !== runId),
{
id: runId,
contents: [{ type: 'text', value: message }],
status: 'loading_packages',
},
]);
},
});

const requiredHandlers = detectRequiredHandlers(codeContent);
for (const handler of requiredHandlers) {
Expand All @@ -173,24 +142,11 @@ export function RunCodeButton({

await currentPyodideInstance.runPythonAsync(codeContent);

const runOutput = await currentPyodideInstance.runPythonAsync(
`sys.stdout.getvalue()`,
);

const runOutputByLines: string[] = runOutput.split('\n');

setConsoleOutputs((outputs) => [
...outputs.filter((output) => output.id !== runId),
{
id: generateUUID(),
contents: runOutputByLines
.filter((line) => line.trim().length)
.map((line) => ({
type: line.startsWith('data:image/png;base64')
? 'image'
: 'text',
value: line,
})),
contents: stdOutputs.filter((output) => output.value.trim().length),
status: 'completed',
},
]);
Expand All @@ -213,6 +169,8 @@ export function RunCodeButton({
try {
pyodide.runPythonAsync(`
import sys
import gc
has_plt = 'matplotlib.pyplot' in sys.modules
if has_plt:
Expand All @@ -238,9 +196,15 @@ export function RunCodeButton({
loadAndRunPython();
});
}}
disabled={block.status === 'streaming'}
disabled={isBlockStreaming}
>
<PlayIcon size={18} /> Run
</Button>
);
}

export const RunCodeButton = memo(PureRunCodeButton, (prevProps, nextProps) => {
if (prevProps.block.status !== nextProps.block.status) return false;

return true;
});

0 comments on commit 1934fe3

Please sign in to comment.