interpreter: load data files with an explicit ?format= override

app/buck2_core/src/bzl.rs

@@ -13,6 +13,7 @@ use std::fmt::Formatter;
 
 use allocative::Allocative;
 use buck2_fs::paths::file_name::FileName;
+use dupe::Dupe;
 use pagable::Pagable;
 use strong_hash::StrongHash;
 
@@ -27,8 +28,64 @@ use crate::cells::paths::CellRelativePath;
 enum ImportPathError {
     #[error("Invalid import path `{0}`")]
     Invalid(CellPath),
-    #[error("Import path must have suffix `.bzl`, `.json`, or `.toml`: `{0}`")]
+    #[error(
+        "Import path must have suffix `.bzl`, `.json`, or `.toml`, or specify `?format=`: `{0}`"
+    )]
     Suffix(CellPath),
+    #[error("Unknown load format `{0}`, expected one of `bzl`, `json`, or `toml`")]
+    UnknownFormat(String),
+}
+
+/// How the contents of a loaded file are parsed.
+#[derive(
+    Clone,
+    Copy,
+    Dupe,
+    Hash,
+    StrongHash,
+    Eq,
+    PartialEq,
+    Debug,
+    Allocative,
+    Pagable
+)]
+pub enum LoadFormat {
+    Bzl,
+    Json,
+    Toml,
+}
+
+impl LoadFormat {
+    pub const fn as_str(self) -> &'static str {
+        match self {
+            LoadFormat::Bzl => "bzl",
+            LoadFormat::Json => "json",
+            LoadFormat::Toml => "toml",
+        }
+    }
+
+    pub fn parse(format: &str) -> buck2_error::Result<Self> {
+        match format {
+            "bzl" => Ok(LoadFormat::Bzl),
+            "json" => Ok(LoadFormat::Json),
+            "toml" => Ok(LoadFormat::Toml),
+            _ => Err(ImportPathError::UnknownFormat(format.to_owned()).into()),
+        }
+    }
+
+    fn from_extension(extension: Option<&str>) -> Self {
+        match extension {
+            Some("json") => LoadFormat::Json,
+            Some("toml") => LoadFormat::Toml,
+            _ => LoadFormat::Bzl,
+        }
+    }
+}
+
+impl Display for LoadFormat {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.as_str())
+    }
 }
 
 /// Path of a `.bzl` file.
@@ -40,6 +97,10 @@ pub struct ImportPath {
     /// The cell of the top-level build module that this is being loaded
     /// (perhaps transitively) into.
     build_file_cell: BuildFileCell,
+    /// Explicit `?format=` from the `load()` string, or `None` to infer from the extension.
+    /// This participates in equality and hashing: the same file loaded under two formats is
+    /// two distinct modules, and both must remain separately cached and invalidated.
+    format: Option<LoadFormat>,
 }
 
 impl ImportPath {
@@ -54,6 +115,17 @@ impl ImportPath {
     pub fn new_with_build_file_cells(
         path: CellPath,
         build_file_cell: BuildFileCell,
+    ) -> buck2_error::Result<Self> {
+        Self::new_with_format(path, build_file_cell, None)
+    }
+
+    /// Like [`ImportPath::new_with_build_file_cells`], but with an explicit `?format=` from the
+    /// `load()` string. An explicit format lifts the extension requirement, so that e.g.
+    /// `Cargo.lock` can be loaded as TOML.
+    pub fn new_with_format(
+        path: CellPath,
+        build_file_cell: BuildFileCell,
+        format: Option<LoadFormat>,
     ) -> buck2_error::Result<Self> {
         if path.parent().is_none() {
             return Err(ImportPathError::Invalid(path).into());
@@ -63,13 +135,29 @@ impl ImportPath {
             return Err(ImportPathError::Invalid(path).into());
         }
 
-        if !matches!(path.path().extension(), Some("bzl" | "json" | "toml")) {
+        let extension_is_known = matches!(path.path().extension(), Some("bzl" | "json" | "toml"));
+        if format.is_none() && !extension_is_known {
             return Err(ImportPathError::Suffix(path).into());
         }
 
+        // A format that merely restates the extension must collapse to `None`, so that
+        // `:a.bzl?format=bzl` and `:a.bzl` are one module rather than two. Evaluating the same
+        // `.bzl` twice would hand out two copies of its providers and transitive sets, and
+        // anything doing pointer equality on those would then silently go wrong.
+        let format = match format {
+            Some(format)
+                if extension_is_known
+                    && format == LoadFormat::from_extension(path.path().extension()) =>
+            {
+                None
+            }
+            format => format,
+        };
+
         Ok(Self {
             path,
             build_file_cell,
+            format,
         })
     }
 
@@ -89,6 +177,7 @@ impl ImportPath {
         Ok(Self {
             path,
             build_file_cell,
+            format: None,
         })
     }
 
@@ -128,6 +217,19 @@ impl ImportPath {
         &self.path
     }
 
+    /// How this file should be parsed: the explicit `?format=`, else inferred from the extension.
+    pub fn load_format(&self) -> LoadFormat {
+        self.format
+            .unwrap_or_else(|| LoadFormat::from_extension(self.path.path().extension()))
+    }
+
+    /// Whether [`ImportPath::load_format`] came from an explicit `?format=` rather than the
+    /// extension. Only useful for diagnostics: the extension is otherwise the reader's only
+    /// clue as to how the file was parsed.
+    pub fn has_explicit_format(&self) -> bool {
+        self.format.is_some()
+    }
+
     /// Parent directory of the import path.
     pub fn path_parent(&self) -> CellPathRef<'_> {
         self.path
@@ -139,9 +241,100 @@ impl ImportPath {
 impl Display for ImportPath {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         if self.build_file_cell.name() == self.path.cell() {
-            write!(f, "{}", self.path)
+            write!(f, "{}", self.path)?;
         } else {
-            write!(f, "{}@{}", self.path, self.build_file_cell.name())
+            write!(f, "{}@{}", self.path, self.build_file_cell.name())?;
+        }
+        if let Some(format) = self.format {
+            write!(f, "?format={format}")?;
         }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn import(filename: &str, format: Option<LoadFormat>) -> buck2_error::Result<ImportPath> {
+        let path = CellPath::new(
+            CellName::testing_new("root"),
+            CellRelativePath::unchecked_new("pkg").join(FileName::unchecked_new(filename)),
+        );
+        ImportPath::new_with_format(
+            path,
+            BuildFileCell::new(CellName::testing_new("root")),
+            format,
+        )
+    }
+
+    #[test]
+    fn format_is_inferred_from_extension() {
+        assert_eq!(LoadFormat::Bzl, import("a.bzl", None).unwrap().load_format());
+        assert_eq!(
+            LoadFormat::Json,
+            import("a.json", None).unwrap().load_format()
+        );
+        assert_eq!(
+            LoadFormat::Toml,
+            import("a.toml", None).unwrap().load_format()
+        );
+    }
+
+    #[test]
+    fn unknown_extension_requires_an_explicit_format() {
+        for filename in ["Cargo.lock", "Makefile", "rules.bzl.in"] {
+            assert!(import(filename, None).is_err(), "{filename}");
+
+            let explicit = import(filename, Some(LoadFormat::Toml)).unwrap();
+            assert_eq!(LoadFormat::Toml, explicit.load_format(), "{filename}");
+            assert!(explicit.has_explicit_format(), "{filename}");
+        }
+    }
+
+    #[test]
+    fn an_extensionless_file_keeps_a_redundant_looking_format() {
+        // `from_extension` defaults to `Bzl`, but an extensionless file has no extension for the
+        // format to be redundant with, so it must stay explicit or the path would not construct.
+        let explicit = import("Makefile", Some(LoadFormat::Bzl)).unwrap();
+        assert_eq!(LoadFormat::Bzl, explicit.load_format());
+        assert!(explicit.has_explicit_format());
+    }
+
+    #[test]
+    fn explicit_format_overrides_the_extension() {
+        let as_toml = import("a.json", Some(LoadFormat::Toml)).unwrap();
+        assert_eq!(LoadFormat::Toml, as_toml.load_format());
+        // The same file under two formats must stay two distinct modules.
+        assert_ne!(import("a.json", None).unwrap(), as_toml);
+    }
+
+    #[test]
+    fn redundant_format_collapses_so_the_module_is_not_duplicated() {
+        for (filename, format) in [
+            ("a.bzl", LoadFormat::Bzl),
+            ("a.json", LoadFormat::Json),
+            ("a.toml", LoadFormat::Toml),
+        ] {
+            let explicit = import(filename, Some(format)).unwrap();
+            assert_eq!(import(filename, None).unwrap(), explicit);
+            assert!(!explicit.has_explicit_format());
+        }
+    }
+
+    #[test]
+    fn display_round_trips_an_explicit_format() {
+        assert_eq!("root//pkg/a.bzl", import("a.bzl", None).unwrap().to_string());
+        assert_eq!(
+            "root//pkg/Cargo.lock?format=toml",
+            import("Cargo.lock", Some(LoadFormat::Toml))
+                .unwrap()
+                .to_string()
+        );
+    }
+
+    #[test]
+    fn a_path_containing_a_question_mark_is_still_rejected() {
+        assert!(import("a.bzl?format=toml", Some(LoadFormat::Toml)).is_err());
     }
 }

app/buck2_interpreter/src/load_module.rs

@@ -10,6 +10,7 @@
 
 use async_trait::async_trait;
 use buck2_core::bzl::ImportPath;
+use buck2_core::bzl::LoadFormat;
 use buck2_core::package::PackageLabel;
 use buck2_util::late_binding::LateBinding;
 use dice::DiceComputations;
@@ -65,10 +66,10 @@ pub trait InterpreterCalculation {
         &mut self,
         path: &ImportPath,
     ) -> buck2_error::Result<LoadedModule> {
-        let module_path = match path.path().path().extension() {
-            Some("json") => StarlarkModulePath::JsonFile(path),
-            Some("toml") => StarlarkModulePath::TomlFile(path),
-            _ => StarlarkModulePath::LoadFile(path),
+        let module_path = match path.load_format() {
+            LoadFormat::Json => StarlarkModulePath::JsonFile(path),
+            LoadFormat::Toml => StarlarkModulePath::TomlFile(path),
+            LoadFormat::Bzl => StarlarkModulePath::LoadFile(path),
         };
         self.get_loaded_module(module_path).await
     }

app/buck2_interpreter/src/parse_import.rs

@@ -11,6 +11,7 @@
 //! Parses imports for load_file() calls in build files.
 
 use buck2_core::bzl::ImportPath;
+use buck2_core::bzl::LoadFormat;
 use buck2_core::cells::CellAliasResolver;
 use buck2_core::cells::build_file_cell::BuildFileCell;
 use buck2_core::cells::cell_path::CellPath;
@@ -37,6 +38,24 @@ enum ImportParseError {
         "Unable to parse import spec. Expected format `(@<cell>)//package/name:filename.bzl` or `:filename.bzl`, but got a path. Got `{0}`"
     )]
     NotAFileName(String),
+    #[error("Unable to parse import spec. Expected a `?format=` query, got `?{1}` in `{0}`")]
+    UnknownQuery(String, String),
+}
+
+/// Split a trailing `?format=<fmt>` query off a `load()` string.
+///
+/// The format overrides the extension-based choice of parser, so that files whose names carry
+/// no usable extension (`Cargo.lock`) or a misleading one (`rules.bzl.in`) can still be loaded.
+pub fn split_load_format(import: &str) -> buck2_error::Result<(&str, Option<LoadFormat>)> {
+    match import.split_once('?') {
+        None => Ok((import, None)),
+        Some((path, query)) => match query.strip_prefix("format=") {
+            Some(format) => Ok((path, Some(LoadFormat::parse(format)?))),
+            None => {
+                Err(ImportParseError::UnknownQuery(import.to_owned(), query.to_owned()).into())
+            }
+        },
+    }
 }
 
 pub enum RelativeImports<'a> {
@@ -421,4 +440,26 @@ mod tests {
         };
         Ok(())
     }
+
+    #[test]
+    fn test_split_load_format() -> buck2_error::Result<()> {
+        assert_eq!((":a.bzl", None), split_load_format(":a.bzl")?);
+        assert_eq!(
+            (":Cargo.lock", Some(LoadFormat::Toml)),
+            split_load_format(":Cargo.lock?format=toml")?
+        );
+        assert_eq!(
+            ("cell//pkg:package.json.tmpl", Some(LoadFormat::Json)),
+            split_load_format("cell//pkg:package.json.tmpl?format=json")?
+        );
+        assert_eq!(
+            (":rules.bzl.in", Some(LoadFormat::Bzl)),
+            split_load_format(":rules.bzl.in?format=bzl")?
+        );
+
+        assert!(split_load_format(":a.bzl?format=yaml").is_err());
+        assert!(split_load_format(":a.bzl?other=1").is_err());
+        assert!(split_load_format(":a.bzl?").is_err());
+        Ok(())
+    }
 }