|
| 1 | +use clap::{arg, Arg, ArgAction, ArgMatches, Command}; |
| 2 | +use miette::{IntoDiagnostic, Result}; |
| 3 | + |
| 4 | +fn build_parser() -> Command { |
| 5 | + Command::new("python_like_parser") |
| 6 | + .disable_version_flag(true) |
| 7 | + .disable_help_flag(true) |
| 8 | + .dont_delimit_trailing_values(true) |
| 9 | + .allow_hyphen_values(true) |
| 10 | + .arg(Arg::new("ignore_env").short('E').action(ArgAction::SetTrue)) |
| 11 | + .arg(Arg::new("isolate").short('I').action(ArgAction::SetTrue)) |
| 12 | + .arg( |
| 13 | + Arg::new("no_user_site") |
| 14 | + .short('s') |
| 15 | + .action(ArgAction::SetTrue), |
| 16 | + ) |
| 17 | + .arg( |
| 18 | + Arg::new("no_import_site") |
| 19 | + .short('S') |
| 20 | + .action(ArgAction::SetTrue), |
| 21 | + ) |
| 22 | + .arg( |
| 23 | + |
| 24 | + .trailing_var_arg(true) |
| 25 | + .required(false) |
| 26 | + .allow_hyphen_values(true), |
| 27 | + ) |
| 28 | +} |
| 29 | + |
| 30 | +pub struct ArgState { |
| 31 | + pub ignore_env: bool, |
| 32 | + pub isolate: bool, |
| 33 | + pub no_import_site: bool, |
| 34 | + pub no_user_site: bool, |
| 35 | + pub remaining_args: Vec<String>, |
| 36 | +} |
| 37 | + |
| 38 | +fn extract_state(matches: &ArgMatches) -> ArgState { |
| 39 | + ArgState { |
| 40 | + // E and I are crucial for transformation |
| 41 | + ignore_env: *matches.get_one::<bool>("ignore_env").unwrap_or(&false), |
| 42 | + isolate: *matches.get_one::<bool>("isolate").unwrap_or(&false), |
| 43 | + |
| 44 | + // s is crucial for transformation |
| 45 | + no_user_site: *matches.get_one::<bool>("no_user_site").unwrap_or(&false), |
| 46 | + |
| 47 | + no_import_site: *matches.get_one::<bool>("no_import_site").unwrap_or(&false), |
| 48 | + |
| 49 | + remaining_args: matches |
| 50 | + .get_many::<String>("args") |
| 51 | + .unwrap_or_default() |
| 52 | + .map(|it| it.to_string()) |
| 53 | + .collect::<Vec<_>>(), |
| 54 | + } |
| 55 | +} |
| 56 | + |
| 57 | +pub fn reparse_args(original_argv: &Vec<&str>) -> Result<Vec<String>> { |
| 58 | + let parser = build_parser(); |
| 59 | + let matches = parser |
| 60 | + .try_get_matches_from(original_argv) |
| 61 | + .into_diagnostic()?; |
| 62 | + let parsed_args = extract_state(&matches); |
| 63 | + |
| 64 | + let mut argv: Vec<String> = Vec::new(); |
| 65 | + let push_flag = |argv: &mut Vec<String>, flag: char, is_set: bool| { |
| 66 | + if is_set { |
| 67 | + argv.push(format!("-{}", flag)); |
| 68 | + } |
| 69 | + }; |
| 70 | + |
| 71 | + // Retain the original argv binary |
| 72 | + argv.push(original_argv[0].to_string()); |
| 73 | + |
| 74 | + // -I replacement logic: -I is never pushed, its effects (-E and -s) are handled separately. |
| 75 | + // -E removal: -E is never pushd |
| 76 | + // -s inclusion logic: we ALWAYS push -s |
| 77 | + push_flag( |
| 78 | + &mut argv, |
| 79 | + 's', |
| 80 | + parsed_args.no_user_site | parsed_args.isolate, |
| 81 | + ); |
| 82 | + |
| 83 | + push_flag(&mut argv, 'S', parsed_args.no_import_site); |
| 84 | + |
| 85 | + argv.extend(parsed_args.remaining_args.iter().cloned()); |
| 86 | + |
| 87 | + Ok(argv) |
| 88 | +} |
| 89 | + |
| 90 | +#[cfg(test)] |
| 91 | +mod test { |
| 92 | + use super::*; |
| 93 | + |
| 94 | + #[test] |
| 95 | + fn basic_args_preserved1() { |
| 96 | + let orig = vec!["python", "-B", "-s", "script.py", "arg1"]; |
| 97 | + let reparsed = reparse_args(&orig); |
| 98 | + assert!(reparsed.is_ok(), "Args failed to parse {:?}", reparsed); |
| 99 | + let reparsed = reparsed.unwrap(); |
| 100 | + assert!( |
| 101 | + orig == reparsed, |
| 102 | + "Args shouldn't have changed, got {:?}", |
| 103 | + reparsed |
| 104 | + ); |
| 105 | + } |
| 106 | + |
| 107 | + #[test] |
| 108 | + fn basic_args_preserved2() { |
| 109 | + let orig = vec!["python", "-s", "-c", "exit(0)", "arg1"]; |
| 110 | + let reparsed = reparse_args(&orig); |
| 111 | + assert!(&reparsed.is_ok(), "Args failed to parse {:?}", reparsed); |
| 112 | + let reparsed = reparsed.unwrap(); |
| 113 | + assert!( |
| 114 | + orig == reparsed, |
| 115 | + "Args shouldn't have changed, got {:?}", |
| 116 | + reparsed |
| 117 | + ); |
| 118 | + } |
| 119 | + |
| 120 | + #[test] |
| 121 | + fn basic_binary_preserved() { |
| 122 | + let orig = vec![ |
| 123 | + "/some/arbitrary/path/python", |
| 124 | + "-B", |
| 125 | + "-s", |
| 126 | + "script.py", |
| 127 | + "arg1", |
| 128 | + ]; |
| 129 | + let reparsed = reparse_args(&orig); |
| 130 | + assert!(reparsed.is_ok(), "Args failed to parse {:?}", reparsed); |
| 131 | + let reparsed = reparsed.unwrap(); |
| 132 | + assert!( |
| 133 | + orig == reparsed, |
| 134 | + "Args shouldn't have changed, got {:?}", |
| 135 | + reparsed |
| 136 | + ); |
| 137 | + } |
| 138 | + |
| 139 | + #[test] |
| 140 | + fn basic_s_remains() { |
| 141 | + // We expect to not modify the -s flag |
| 142 | + let orig = vec!["python", "-s", "-c", "exit(0)", "arg1"]; |
| 143 | + let reparsed = reparse_args(&orig); |
| 144 | + assert!(reparsed.is_ok(), "Args failed to parse {:?}", reparsed); |
| 145 | + assert!(orig == reparsed.unwrap(), "Args shouldn't have changed"); |
| 146 | + } |
| 147 | + |
| 148 | + #[test] |
| 149 | + fn basic_e_gets_unset() { |
| 150 | + // We expect to REMOVE the -E flag |
| 151 | + let orig = vec!["python", "-E", "-s", "-c", "exit(0)", "arg1"]; |
| 152 | + let expected = vec!["python", "-s", "-c", "exit(0)", "arg1"]; |
| 153 | + let reparsed = reparse_args(&orig); |
| 154 | + assert!(reparsed.is_ok(), "Args failed to parse {:?}", reparsed); |
| 155 | + assert!(expected == reparsed.unwrap(), "-E wasn't unset"); |
| 156 | + } |
| 157 | + |
| 158 | + #[test] |
| 159 | + fn basic_i_becomes_s() { |
| 160 | + // We expect to CONVERT the -I flag to -E (which we ignore) and -s (which we keep) |
| 161 | + let orig = vec!["python", "-I", "-c", "exit(0)", "arg1"]; |
| 162 | + let expected = vec!["python", "-s", "-c", "exit(0)", "arg1"]; |
| 163 | + let reparsed = reparse_args(&orig); |
| 164 | + assert!(reparsed.is_ok(), "Args failed to parse {:?}", reparsed); |
| 165 | + assert!(expected == reparsed.unwrap(), "Didn't translate -I to -s"); |
| 166 | + } |
| 167 | + |
| 168 | + #[test] |
| 169 | + fn basic_m_preserved() { |
| 170 | + // We expect to ADD the -s flag |
| 171 | + let orig = vec!["python", "-m", "build", "--unknown", "arg1"]; |
| 172 | + let expected = vec!["python", "-m", "build", "--unknown", "arg1"]; |
| 173 | + let reparsed = reparse_args(&orig); |
| 174 | + assert!(reparsed.is_ok(), "Args failed to parse {:?}", reparsed); |
| 175 | + assert!(expected == reparsed.unwrap(), "Didn't add -s"); |
| 176 | + } |
| 177 | + |
| 178 | + #[test] |
| 179 | + fn basic_trailing_args_preserved() { |
| 180 | + let orig = vec![ |
| 181 | + "python3", |
| 182 | + "uv/private/sdist_build/build_helper.py", |
| 183 | + "bazel-out/darwin_arm64-fastbuild/bin/external/+uv+sbuild__pypi__default__bravado_core/src", |
| 184 | + "bazel-out/darwin_arm64-fastbuild/bin/external/+uv+sbuild__pypi__default__bravado_core/build" |
| 185 | + ]; |
| 186 | + let expected = vec![ |
| 187 | + "python3", |
| 188 | + "uv/private/sdist_build/build_helper.py", |
| 189 | + "bazel-out/darwin_arm64-fastbuild/bin/external/+uv+sbuild__pypi__default__bravado_core/src", |
| 190 | + "bazel-out/darwin_arm64-fastbuild/bin/external/+uv+sbuild__pypi__default__bravado_core/build" |
| 191 | + ]; |
| 192 | + let reparsed = reparse_args(&orig); |
| 193 | + assert!(reparsed.is_ok(), "Args failed to parse {:?}", reparsed); |
| 194 | + let reparsed = reparsed.unwrap(); |
| 195 | + assert!( |
| 196 | + expected == reparsed, |
| 197 | + "Something happened to the args, got {:?}", |
| 198 | + reparsed |
| 199 | + ); |
| 200 | + } |
| 201 | + |
| 202 | + #[test] |
| 203 | + fn m_preserved_s_added_varargs_preserved() { |
| 204 | + let orig = vec![ |
| 205 | + "python3", |
| 206 | + "-m", |
| 207 | + "build", |
| 208 | + "--no-isolation", |
| 209 | + "--out-dir", |
| 210 | + "bazel-out/darwin_arm64-fastbuild/bin/external/+uv+sbuild__pypi__default__bravado_core/build", |
| 211 | + "bazel-out/darwin_arm64-fastbuild/bin/external/+uv+sbuild__pypi__default__bravado_core/src", |
| 212 | + ]; |
| 213 | + let expected = vec![ |
| 214 | + "python3", |
| 215 | + "-m", |
| 216 | + "build", |
| 217 | + "--no-isolation", |
| 218 | + "--out-dir", |
| 219 | + "bazel-out/darwin_arm64-fastbuild/bin/external/+uv+sbuild__pypi__default__bravado_core/build", |
| 220 | + "bazel-out/darwin_arm64-fastbuild/bin/external/+uv+sbuild__pypi__default__bravado_core/src", |
| 221 | + ]; |
| 222 | + let reparsed = reparse_args(&orig); |
| 223 | + assert!(reparsed.is_ok(), "Args failed to parse {:?}", reparsed); |
| 224 | + let reparsed = reparsed.unwrap(); |
| 225 | + assert!( |
| 226 | + expected == reparsed, |
| 227 | + "Something happened to the args, got {:?}", |
| 228 | + reparsed |
| 229 | + ); |
| 230 | + } |
| 231 | +} |
0 commit comments