Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.Collections;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.Map;
Expand Down Expand Up @@ -76,6 +77,18 @@ protected boolean removeEldestEntry(Map.Entry<String, Path> eldest) {
* @return Path to the workspace directory containing .venv
*/
static Path getOrCreateWorkspace(String pyprojectContent) {
return getOrCreateWorkspace(pyprojectContent, Collections.emptyMap());
}

/**
* Gets or creates a workspace directory for the given pyproject.toml content.
* Workspaces are cached by content hash to avoid repeated installations.
*
* @param pyprojectContent The complete pyproject.toml file content
* @param environment additional environment variables for subprocess (e.g., SSL_CERT_FILE)
* @return Path to the workspace directory containing .venv
*/
static Path getOrCreateWorkspace(String pyprojectContent, Map<String, String> environment) {
String hash = hashContent(pyprojectContent);

// Check in-memory cache
Expand Down Expand Up @@ -106,10 +119,10 @@ static Path getOrCreateWorkspace(String pyprojectContent) {
);

// Sync: creates .venv, generates uv.lock, and installs dependencies
runCommand(tempDir, "uv", "sync");
runCommand(tempDir, environment, "uv", "sync");

// Install ty for type stubs
runCommand(tempDir, "uv", "pip", "install", "ty");
runCommand(tempDir, environment, "uv", "pip", "install", "ty");

// Write workspace version for cache invalidation
Files.write(tempDir.resolve("version.txt"),
Expand Down Expand Up @@ -152,6 +165,21 @@ static Path getOrCreateWorkspace(String pyprojectContent) {
*/
static @Nullable Path getOrCreateRequirementsWorkspace(String requirementsContent,
@Nullable Path originalFilePath) {
return getOrCreateRequirementsWorkspace(requirementsContent, originalFilePath, Collections.emptyMap());
}

/**
* Gets or creates a workspace directory for a requirements.txt file.
* Returns null (graceful degradation) when uv is unavailable.
*
* @param requirementsContent The complete requirements.txt content
* @param originalFilePath The original file path on disk (supports -r includes), or null
* @param environment additional environment variables for subprocess (e.g., SSL_CERT_FILE)
* @return Path to the workspace directory, or null if uv is unavailable
*/
static @Nullable Path getOrCreateRequirementsWorkspace(String requirementsContent,
@Nullable Path originalFilePath,
Map<String, String> environment) {
String uvPath = UvExecutor.findUvExecutable();
if (uvPath == null) {
return null;
Expand Down Expand Up @@ -180,7 +208,7 @@ static Path getOrCreateWorkspace(String pyprojectContent) {

try {
// Create virtualenv
runCommandWithPath(tempDir, uvPath, "venv");
runCommandWithPath(tempDir, uvPath, environment, "venv");

// Install dependencies from requirements file
Path reqFile;
Expand All @@ -190,10 +218,10 @@ static Path getOrCreateWorkspace(String pyprojectContent) {
reqFile = tempDir.resolve("requirements.txt");
Files.write(reqFile, requirementsContent.getBytes(StandardCharsets.UTF_8));
}
runCommandWithPath(tempDir, uvPath, "pip", "install", "-r", reqFile.toString());
runCommandWithPath(tempDir, uvPath, environment, "pip", "install", "-r", reqFile.toString());

// Capture freeze output BEFORE installing ty
UvExecutor.RunResult freezeResult = UvExecutor.run(tempDir, uvPath, "pip", "freeze");
UvExecutor.RunResult freezeResult = UvExecutor.run(tempDir, uvPath, environment, "pip", "freeze");
if (freezeResult.isSuccess()) {
Files.write(
tempDir.resolve("freeze.txt"),
Expand All @@ -202,7 +230,7 @@ static Path getOrCreateWorkspace(String pyprojectContent) {
}

// Install ty for type stubs (after freeze so it's not in the dep model)
runCommandWithPath(tempDir, uvPath, "pip", "install", "ty");
runCommandWithPath(tempDir, uvPath, environment, "pip", "install", "ty");

// Write workspace version for cache invalidation
Files.write(tempDir.resolve("version.txt"),
Expand Down Expand Up @@ -243,6 +271,22 @@ static Path getOrCreateWorkspace(String pyprojectContent) {
*/
public static @Nullable Path getOrCreateSetuptoolsWorkspace(String manifestContent,
@Nullable Path projectDir) {
return getOrCreateSetuptoolsWorkspace(manifestContent, projectDir, Collections.emptyMap());
}

/**
* Gets or creates a workspace directory for a setuptools project (setup.cfg / setup.py).
* Uses {@code uv pip install <projectDir>} to install the project and its dependencies.
* Returns null (graceful degradation) when uv is unavailable.
*
* @param manifestContent The setup.cfg (or setup.py) content for hashing
* @param projectDir The project directory to install from, or null
* @param environment additional environment variables for subprocess (e.g., SSL_CERT_FILE)
* @return Path to the workspace directory, or null if uv is unavailable
*/
public static @Nullable Path getOrCreateSetuptoolsWorkspace(String manifestContent,
@Nullable Path projectDir,
Map<String, String> environment) {
String uvPath = UvExecutor.findUvExecutable();
if (uvPath == null) {
return null;
Expand Down Expand Up @@ -275,13 +319,13 @@ static Path getOrCreateWorkspace(String pyprojectContent) {

try {
// Create virtualenv
runCommandWithPath(tempDir, uvPath, "venv");
runCommandWithPath(tempDir, uvPath, environment, "venv");

// Install from the project directory
runCommandWithPath(tempDir, uvPath, "pip", "install", projectDir.toString());
runCommandWithPath(tempDir, uvPath, environment, "pip", "install", projectDir.toString());

// Capture freeze output BEFORE installing ty
UvExecutor.RunResult freezeResult = UvExecutor.run(tempDir, uvPath, "pip", "freeze");
UvExecutor.RunResult freezeResult = UvExecutor.run(tempDir, uvPath, environment, "pip", "freeze");
if (freezeResult.isSuccess()) {
Files.write(
tempDir.resolve("freeze.txt"),
Expand All @@ -290,7 +334,7 @@ static Path getOrCreateWorkspace(String pyprojectContent) {
}

// Install ty for type stubs (after freeze so it's not in the dep model)
runCommandWithPath(tempDir, uvPath, "pip", "install", "ty");
runCommandWithPath(tempDir, uvPath, environment, "pip", "install", "ty");

// Write workspace version for cache invalidation
Files.write(tempDir.resolve("version.txt"),
Expand Down Expand Up @@ -335,14 +379,14 @@ private static boolean isRequirementsWorkspaceValid(Path workspaceDir) {
hasCurrentVersion(workspaceDir);
}

private static void runCommandWithPath(Path dir, String uvPath, String... args) throws IOException, InterruptedException {
UvExecutor.RunResult result = UvExecutor.run(dir, uvPath, args);
private static void runCommandWithPath(Path dir, String uvPath, Map<String, String> environment, String... args) throws IOException, InterruptedException {
UvExecutor.RunResult result = UvExecutor.run(dir, uvPath, environment, args);
if (!result.isSuccess()) {
throw new RuntimeException("uv " + String.join(" ", args) + " failed with exit code: " + result.getExitCode());
}
}

private static void runCommand(Path dir, String... command) throws IOException, InterruptedException {
private static void runCommand(Path dir, Map<String, String> environment, String... command) throws IOException, InterruptedException {
String uvPath = UvExecutor.findUvExecutable();
if (uvPath == null) {
throw new RuntimeException("uv is not installed. Install it with: pip install uv");
Expand All @@ -352,7 +396,7 @@ private static void runCommand(Path dir, String... command) throws IOException,
String[] args = new String[command.length - 1];
System.arraycopy(command, 1, args, 0, args.length);

UvExecutor.RunResult result = UvExecutor.run(dir, uvPath, args);
UvExecutor.RunResult result = UvExecutor.run(dir, uvPath, environment, args);
if (!result.isSuccess()) {
throw new RuntimeException(String.join(" ", command) + " failed with exit code: " + result.getExitCode());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@
import java.util.regex.Pattern;
import java.util.stream.Stream;

import static java.util.Collections.emptyMap;

import static org.openrewrite.Tree.randomId;

/**
Expand All @@ -50,6 +52,15 @@ public class RequirementsTxtParser implements Parser {
private static final Pattern EXTRA_MARKER_PATTERN = Pattern.compile("\\bextra\\s*==");

private final PlainTextParser plainTextParser = new PlainTextParser();
private final Map<String, String> subprocessEnvironment;

public RequirementsTxtParser() {
this(emptyMap());
}

public RequirementsTxtParser(Map<String, String> subprocessEnvironment) {
this.subprocessEnvironment = subprocessEnvironment;
}

@Override
public Stream<SourceFile> parseInputs(Iterable<Input> sources, @Nullable Path relativeTo, ExecutionContext ctx) {
Expand All @@ -64,7 +75,7 @@ public Stream<SourceFile> parseInputs(Iterable<Input> sources, @Nullable Path re
originalFilePath = relativeTo.resolve(text.getSourcePath());
}
Path workspace = DependencyWorkspace.getOrCreateRequirementsWorkspace(
text.getText(), originalFilePath);
text.getText(), originalFilePath, subprocessEnvironment);
if (workspace == null) {
return sf;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import java.nio.file.Path;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

import static org.openrewrite.Tree.randomId;
Expand All @@ -41,6 +42,15 @@
public class SetupCfgParser implements Parser {

private final PlainTextParser plainTextParser = new PlainTextParser();
private final Map<String, String> subprocessEnvironment;

public SetupCfgParser() {
this(Collections.emptyMap());
}

public SetupCfgParser(Map<String, String> subprocessEnvironment) {
this.subprocessEnvironment = subprocessEnvironment;
}

@Override
public Stream<SourceFile> parseInputs(Iterable<Input> sources, @Nullable Path relativeTo, ExecutionContext ctx) {
Expand All @@ -56,7 +66,7 @@ public Stream<SourceFile> parseInputs(Iterable<Input> sources, @Nullable Path re
projectDir = filePath.getParent();
}
Path workspace = DependencyWorkspace.getOrCreateSetuptoolsWorkspace(
text.getText(), projectDir);
text.getText(), projectDir, subprocessEnvironment);
if (workspace == null) {
return sf;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
Expand All @@ -37,7 +38,6 @@ public class UvExecutor {

private static final long DEFAULT_TIMEOUT_SECONDS = 120;
private static @Nullable String cachedUvPath;

@Value
public static class RunResult {
boolean success;
Expand All @@ -47,20 +47,22 @@ public static class RunResult {
}

/**
* Run a uv command in the given directory.
* Run a uv command in the given directory with additional environment variables.
*
* @param workDir the working directory
* @param uvPath the path to the uv executable
* @param args the arguments to pass to uv
* @param workDir the working directory
* @param uvPath the path to the uv executable
* @param environment additional environment variables (e.g., SSL_CERT_FILE, HTTP_PROXY)
* @param args the arguments to pass to uv
* @return the run result
*/
public static RunResult run(Path workDir, String uvPath, String... args) throws IOException, InterruptedException {
public static RunResult run(Path workDir, String uvPath, Map<String, String> environment, String... args) throws IOException, InterruptedException {
String[] command = new String[args.length + 1];
command[0] = uvPath;
System.arraycopy(args, 0, command, 1, args.length);

ProcessBuilder pb = new ProcessBuilder(command);
pb.directory(workDir.toFile());
pb.environment().putAll(environment);
pb.redirectErrorStream(false);

Process process = pb.start();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,9 @@
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Collections;
import java.util.Comparator;
import java.util.Map;

/**
* Utility for regenerating uv.lock files by running {@code uv lock} in a temporary directory.
Expand Down Expand Up @@ -53,7 +55,7 @@ public static Result failure(String errorMessage) {
* @return a result containing the new lock file content, or an error message
*/
public static Result regenerate(String pyprojectContent) {
return regenerate(pyprojectContent, null);
return regenerate(pyprojectContent, null, Collections.emptyMap());
}

/**
Expand All @@ -67,6 +69,21 @@ public static Result regenerate(String pyprojectContent) {
* @return a result containing the new lock file content, or an error message
*/
public static Result regenerate(String pyprojectContent, @Nullable String existingLockContent) {
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same comment as above. AIs tends to want to leave things for backwards compatibility, but if leaving it opens us up for bugs, better to remove aggressively.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The other one was unused so I removed it, this one is still being used.

return regenerate(pyprojectContent, existingLockContent, Collections.emptyMap());
}

/**
* Regenerate a uv.lock file from the given pyproject.toml content.
* When an existing lock file is provided it is seeded into the working
* directory so that {@code uv lock} performs a minimal update rather
* than re-resolving every dependency from scratch.
*
* @param pyprojectContent the pyproject.toml content to lock
* @param existingLockContent the current uv.lock content, or {@code null}
* @param environment additional environment variables for subprocess (e.g., SSL_CERT_FILE)
* @return a result containing the new lock file content, or an error message
*/
public static Result regenerate(String pyprojectContent, @Nullable String existingLockContent, Map<String, String> environment) {
String uvPath = UvExecutor.findUvExecutable();
if (uvPath == null) {
return Result.failure("uv is not installed. Install it with: pip install uv");
Expand All @@ -88,7 +105,7 @@ public static Result regenerate(String pyprojectContent, @Nullable String existi
);
}

UvExecutor.RunResult runResult = UvExecutor.run(tempDir, uvPath, "lock");
UvExecutor.RunResult runResult = UvExecutor.run(tempDir, uvPath, environment, "lock");
if (!runResult.isSuccess()) {
return Result.failure("uv lock failed (exit code " + runResult.getExitCode() + "): " + runResult.getStderr());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -273,7 +273,7 @@ public int characteristics() {
return null;
}

Path workspace = DependencyWorkspace.getOrCreateSetuptoolsWorkspace(source, projectPath);
Path workspace = DependencyWorkspace.getOrCreateSetuptoolsWorkspace(source, projectPath, commandEnv);
if (workspace == null) {
return null;
}
Expand Down Expand Up @@ -334,11 +334,11 @@ private Stream<SourceFile> parseManifest(Path projectPath, @Nullable Path relati
Path setupCfgPath = projectPath.resolve("setup.cfg");
if (Files.exists(setupCfgPath)) {
Parser.Input input = Parser.Input.fromFile(setupCfgPath);
return new SetupCfgParser().parseInputs(
return new SetupCfgParser(commandEnv).parseInputs(
Collections.singletonList(input), effectiveRelativeTo, ctx);
}

RequirementsTxtParser reqsParser = new RequirementsTxtParser();
RequirementsTxtParser reqsParser = new RequirementsTxtParser(commandEnv);
try (Stream<Path> entries = Files.list(projectPath)) {
Path reqsPath = entries
.filter(p -> reqsParser.accept(p.getFileName()))
Expand Down Expand Up @@ -719,6 +719,7 @@ private void bootstrapOpenrewrite(Path pipPackagesPath, String version) {
"--target=" + pipPackagesPath.toAbsolutePath().normalize(),
"openrewrite==" + version
);
pb.environment().putAll(environment);
pb.redirectErrorStream(true);
if (log != null) {
File logFile = log.toAbsolutePath().normalize().toFile();
Expand Down