python-tools

import { randomUUID } from "node:crypto";
import { mkdir, readFile, rm, stat, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import path from "node:path";

import {
  resolvePythonCommand,
  runResolvedPythonCommand,
} from "./pythonResolver";
import { ensurePythonFile, MAX_FILE_READ_BYTES } from "./safePaths";

export interface BugCheckIssue {
  severity: "error" | "warning";
  message: string;
  line?: number;
  column?: number;
}

export interface CheckPythonCodeForBugsResult {
  success: boolean;
  pythonExecutableUsed: string;
  syntaxOk: boolean;
  stdout: string;
  stderr: string;
  exitCode: number | null;
  issues: BugCheckIssue[];
}

export interface CheckPythonFileForBugsResult extends CheckPythonCodeForBugsResult {
  filePath: string;
}

const DEFAULT_PY_COMPILE_TIMEOUT_SECONDS = 30;
const MAX_CODE_LENGTH = 50_000;

export async function checkPythonCodeForBugs(input: {
  code: string;
  runPyCompile: boolean;
}): Promise<CheckPythonCodeForBugsResult> {
  if (input.code.trim().length === 0) {
    throw new Error("code cannot be empty.");
  }

  if (input.code.length > MAX_CODE_LENGTH) {
    throw new Error(`code must be ${MAX_CODE_LENGTH} characters or fewer.`);
  }

  const pythonCommand = await resolvePythonCommand();
  const issues: BugCheckIssue[] = collectTextWarnings(input.code);

  if (!input.runPyCompile) {
    return {
      success: true,
      pythonExecutableUsed: pythonCommand.pythonExecutableUsed,
      syntaxOk: true,
      stdout: "",
      stderr: "",
      exitCode: 0,
      issues,
    };
  }

  const temporaryScriptPath = await writeTemporaryPythonFile(input.code);

  try {
    const compileResult = await runResolvedPythonCommand(
      pythonCommand,
      ["-m", "py_compile", temporaryScriptPath],
      DEFAULT_PY_COMPILE_TIMEOUT_SECONDS,
    );

    const syntaxIssue = parsePyCompileError(compileResult.stderr);
    if (syntaxIssue !== null) {
      issues.push(syntaxIssue);
    }

    const syntaxOk = !compileResult.timedOut && compileResult.exitCode === 0;

    return {
      success: syntaxOk && issues.every((issue) => issue.severity !== "error"),
      pythonExecutableUsed: pythonCommand.pythonExecutableUsed,
      syntaxOk,
      stdout: compileResult.stdout,
      stderr: compileResult.stderr,
      exitCode: compileResult.exitCode,
      issues,
    };
  } finally {
    await rm(temporaryScriptPath, { force: true });
  }
}

export async function checkPythonFileForBugs(filePath: string): Promise<CheckPythonFileForBugsResult> {
  const resolvedFilePath = await ensurePythonFile(filePath);
  const fileStats = await stat(resolvedFilePath);

  if (fileStats.size > MAX_FILE_READ_BYTES) {
    throw new Error(`File is too large for bug checking (${fileStats.size} bytes). Maximum is ${MAX_FILE_READ_BYTES} bytes.`);
  }

  const sourceCode = await readFile(resolvedFilePath, "utf8");
  const issues: BugCheckIssue[] = collectTextWarnings(sourceCode);
  const pythonCommand = await resolvePythonCommand();

  const compileResult = await runResolvedPythonCommand(
    pythonCommand,
    ["-m", "py_compile", resolvedFilePath],
    DEFAULT_PY_COMPILE_TIMEOUT_SECONDS,
  );

  const syntaxIssue = parsePyCompileError(compileResult.stderr);
  if (syntaxIssue !== null) {
    issues.push(syntaxIssue);
  }

  const syntaxOk = !compileResult.timedOut && compileResult.exitCode === 0;

  return {
    success: syntaxOk && issues.every((issue) => issue.severity !== "error"),
    filePath: resolvedFilePath,
    pythonExecutableUsed: pythonCommand.pythonExecutableUsed,
    syntaxOk,
    stdout: compileResult.stdout,
    stderr: compileResult.stderr,
    exitCode: compileResult.exitCode,
    issues,
  };
}

function collectTextWarnings(code: string): BugCheckIssue[] {
  const lines = code.replace(/\r\n/g, "\n").split("\n");
  let sawTabIndent = false;
  let sawSpaceIndent = false;
  let tabIndentLine: number | undefined;

  for (let index = 0; index < lines.length; index += 1) {
    const line = lines[index];
    if (line === undefined || line.length === 0) {
      continue;
    }

    const leadingWhitespace = line.match(/^[\t ]+/)?.[0] ?? "";
    if (leadingWhitespace.includes("\t")) {
      sawTabIndent = true;
      tabIndentLine ??= index + 1;
    }

    if (leadingWhitespace.includes(" ")) {
      sawSpaceIndent = true;
    }
  }

  if (sawTabIndent && sawSpaceIndent) {
    return [{
      severity: "warning",
      message: "Code appears to mix tabs and spaces for indentation.",
      ...(tabIndentLine !== undefined ? { line: tabIndentLine } : {}),
    }];
  }

  return [];
}

function parsePyCompileError(stderr: string): BugCheckIssue | null {
  const trimmedStderr = stderr.trim();
  if (trimmedStderr.length === 0) {
    return null;
  }

  const syntaxLine = trimmedStderr
    .split(/\r?\n/)
    .find((line) => line.toLowerCase().includes("syntaxerror"));

  if (syntaxLine === undefined) {
    return {
      severity: "error",
      message: "py_compile reported an error.",
    };
  }

  const lineNumberMatch = trimmedStderr.match(/line\s+(\d+)/i);
  const caretLine = trimmedStderr.split(/\r?\n/).find((line) => line.includes("^"));
  const caretColumn = caretLine === undefined ? undefined : Math.max(1, caretLine.indexOf("^") + 1);
  const lineNumberText = lineNumberMatch?.[1];
  const parsedLineNumber = lineNumberText === undefined ? undefined : Number.parseInt(lineNumberText, 10);

  const issue: BugCheckIssue = {
    severity: "error",
    message: syntaxLine.trim(),
  };

  if (parsedLineNumber !== undefined && Number.isInteger(parsedLineNumber)) {
    issue.line = parsedLineNumber;
  }

  if (caretColumn !== undefined) {
    issue.column = caretColumn;
  }

  return issue;
}

async function writeTemporaryPythonFile(code: string): Promise<string> {
  const temporaryDirectory = path.join(tmpdir(), "lmstudio-python-tools");
  const filePath = path.join(temporaryDirectory, `check-${randomUUID()}.py`);

  await mkdir(temporaryDirectory, { recursive: true });
  await writeFile(filePath, code, "utf8");
  return filePath;
}