Support load("uv.lock?as=toml", "value") syntax

app/buck2_core/src/bzl.rs

@@ -73,6 +73,34 @@ impl ImportPath {
         })
     }
 
+    /// Create an ImportPath for a file with an explicit format override (e.g. `?as=toml`).
+    /// Skips the file extension validation since the format is determined by the hint,
+    /// not the extension.
+    ///
+    /// Do not include the ?as=toml hint as part of the path.
+    pub fn new_with_explicit_format(
+        path: CellPath,
+        build_file_cell: BuildFileCell,
+    ) -> buck2_error::Result<Self> {
+        if path.parent().is_none() {
+            return Err(ImportPathError::Invalid(path).into());
+        }
+
+        if path.path().as_str().contains('?') {
+            return Err(ImportPathError::Invalid(path).into());
+        }
+
+        Ok(Self {
+            path,
+            build_file_cell,
+        })
+    }
+
+    pub fn new_same_cell_with_explicit_format(path: CellPath) -> buck2_error::Result<Self> {
+        let build_file_cell = BuildFileCell::new(path.cell());
+        Self::new_with_explicit_format(path, build_file_cell)
+    }
+
     /// LSP creates imports for non-bzl files.
     pub fn new_hack_for_lsp(
         path: CellPath,

app/buck2_interpreter/src/parse_import.rs

@@ -20,6 +20,27 @@ use buck2_core::cells::paths::CellRelativePathBuf;
 use buck2_fs::paths::RelativePath;
 use buck2_fs::paths::file_name::FileName;
 
+/// Format hint parsed from `?as=toml` or `?as=json` suffix in load() strings.
+/// Allows loading a file as a specific format regardless of its extension,
+/// e.g. `load(":blah.lock?as=toml", "value")`.
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub enum FormatHint {
+    Json,
+    Toml,
+}
+
+/// Strip a `?as=toml` or `?as=json` format hint from a load string.
+/// Returns the stripped string and the format hint if present.
+pub fn strip_format_hint(import: &str) -> (&str, Option<FormatHint>) {
+    if let Some(base) = import.strip_suffix("?as=toml") {
+        (base, Some(FormatHint::Toml))
+    } else if let Some(base) = import.strip_suffix("?as=json") {
+        (base, Some(FormatHint::Json))
+    } else {
+        (import, None)
+    }
+}
+
 #[derive(buck2_error::Error, Debug)]
 #[buck2(input)]
 enum ImportParseError {

app/buck2_interpreter_for_build/src/interpreter/interpreter_for_dir.rs

@@ -40,8 +40,10 @@ use buck2_interpreter::file_loader::LoadedModules;
 use buck2_interpreter::file_type::StarlarkFileType;
 use buck2_interpreter::import_paths::ImplicitImportPaths;
 use buck2_interpreter::package_imports::ImplicitImport;
+use buck2_interpreter::parse_import::FormatHint;
 use buck2_interpreter::parse_import::RelativeImports;
 use buck2_interpreter::parse_import::parse_import;
+use buck2_interpreter::parse_import::strip_format_hint;
 use buck2_interpreter::paths::module::OwnedStarlarkModulePath;
 use buck2_interpreter::paths::module::StarlarkModulePath;
 use buck2_interpreter::paths::package::PackageFilePath;
@@ -189,13 +191,14 @@ impl LoadResolver for InterpreterLoadResolver {
         path: &str,
         location: Option<&FileSpan>,
     ) -> buck2_error::Result<OwnedStarlarkModulePath> {
+        let (path_str, format_hint) = strip_format_hint(path);
         let relative_import_option = RelativeImports::Allow {
             current_dir_with_allowed_relative: &self.config.current_dir_with_allowed_relative_dirs,
         };
         let path = parse_import(
             self.config.cell_info.cell_alias_resolver(),
             relative_import_option,
-            path,
+            path_str,
         )?;
 
         // check for bxl files first before checking for prelude.
@@ -247,23 +250,35 @@ impl LoadResolver for InterpreterLoadResolver {
         // checks in t-sets, which would fail if we had > 1 copy of the prelude.
         if let Some(prelude_import) = self.config.global_state.configuror.prelude_import() {
             if prelude_import.is_prelude_path(&path) {
-                if path.path().extension() == Some("json") {
-                    return Ok(OwnedStarlarkModulePath::JsonFile(
-                        ImportPath::new_same_cell(path)?,
-                    ));
+                let import_path = if format_hint.is_some() {
+                    ImportPath::new_same_cell_with_explicit_format(path)?
                 } else {
-                    return Ok(OwnedStarlarkModulePath::LoadFile(
-                        ImportPath::new_same_cell(path)?,
-                    ));
-                }
+                    ImportPath::new_same_cell(path)?
+                };
+                return Ok(module_path_with_format_hint(import_path, format_hint));
             }
         }
-        let import_path = ImportPath::new_with_build_file_cells(path, self.build_file_cell)?;
-        Ok(match import_path.path().path().extension() {
+        let import_path = if format_hint.is_some() {
+            ImportPath::new_with_explicit_format(path, self.build_file_cell)?
+        } else {
+            ImportPath::new_with_build_file_cells(path, self.build_file_cell)?
+        };
+        Ok(module_path_with_format_hint(import_path, format_hint))
+    }
+}
+
+fn module_path_with_format_hint(
+    import_path: ImportPath,
+    format_hint: Option<FormatHint>,
+) -> OwnedStarlarkModulePath {
+    match format_hint {
+        Some(FormatHint::Json) => OwnedStarlarkModulePath::JsonFile(import_path),
+        Some(FormatHint::Toml) => OwnedStarlarkModulePath::TomlFile(import_path),
+        None => match import_path.path().path().extension() {
             Some("json") => OwnedStarlarkModulePath::JsonFile(import_path),
             Some("toml") => OwnedStarlarkModulePath::TomlFile(import_path),
             _ => OwnedStarlarkModulePath::LoadFile(import_path),
-        })
+        },
     }
 }