Support for @unstable WIT features (#110)

* Add test for current behavior

* Allow for binding generation of multiple unstable features

* Test features are properly enabled via componentize

Author: benbrandt

Cargo.toml

@@ -51,6 +51,7 @@ semver = "1.0.23"
 [dev-dependencies]
 proptest = "1.5.0"
 hex = "0.4.3"
+tempfile = "3.13.0"
 
 [build-dependencies]
 anyhow = "1.0.89"

src/command.rs

@@ -22,7 +22,7 @@ pub struct Options {
     pub command: Command,
 }
 
-#[derive(clap::Args, Debug)]
+#[derive(clap::Args, Clone, Debug)]
 pub struct Common {
     /// File or directory containing WIT document(s)
     #[arg(short = 'd', long)]
@@ -35,6 +35,19 @@ pub struct Common {
     /// Disable non-error output
     #[arg(short = 'q', long)]
     pub quiet: bool,
+
+    /// Comma-separated list of features that should be enabled when processing
+    /// WIT files.
+    ///
+    /// This enables using `@unstable` annotations in WIT files.
+    #[clap(long)]
+    features: Vec<String>,
+
+    /// Whether or not to activate all WIT features when processing WIT files.
+    ///
+    /// This enables using `@unstable` annotations in WIT files.
+    #[clap(long)]
+    all_features: bool,
 }
 
 #[derive(clap::Subcommand, Debug)]
@@ -120,6 +133,8 @@ fn generate_bindings(common: Common, bindings: Bindings) -> Result<()> {
             .wit_path
             .unwrap_or_else(|| Path::new("wit").to_owned()),
         common.world.as_deref(),
+        &common.features,
+        common.all_features,
         bindings.world_module.as_deref(),
         &bindings.output_dir,
     )
@@ -140,6 +155,8 @@ fn componentize(common: Common, componentize: Componentize) -> Result<()> {
     Runtime::new()?.block_on(crate::componentize(
         common.wit_path.as_deref(),
         common.world.as_deref(),
+        &common.features,
+        common.all_features,
         &python_path.iter().map(|s| s.as_str()).collect::<Vec<_>>(),
         &componentize
             .module_worlds
@@ -241,3 +258,156 @@ fn find_dir(name: &str, path: &Path) -> Result<Option<PathBuf>> {
 
     Ok(None)
 }
+
+#[cfg(test)]
+mod tests {
+    use std::io::Write;
+
+    use super::*;
+
+    /// Generates a WIT file which has unstable feature "x"
+    fn gated_x_wit_file() -> Result<tempfile::NamedTempFile, anyhow::Error> {
+        let mut wit = tempfile::Builder::new()
+            .prefix("gated")
+            .suffix(".wit")
+            .tempfile()?;
+        write!(
+            wit,
+            r#"
+            package foo:bar@1.2.3;
+
+            world bindings {{
+                @unstable(feature = x)
+                import x: func();
+                @since(version = 1.2.3)
+                export y: func();
+            }}
+        "#,
+        )?;
+        Ok(wit)
+    }
+
+    #[test]
+    fn unstable_bindings_not_generated() -> Result<()> {
+        // Given a WIT file with gated features
+        let wit = gated_x_wit_file()?;
+        let out_dir = tempfile::tempdir()?;
+
+        // When generating the bindings for this WIT world
+        let common = Common {
+            wit_path: Some(wit.path().into()),
+            world: None,
+            quiet: false,
+            features: vec![],
+            all_features: false,
+        };
+        let bindings = Bindings {
+            output_dir: out_dir.path().into(),
+            world_module: None,
+        };
+        generate_bindings(common, bindings)?;
+
+        // Then the gated feature doesn't appear
+        let generated = fs::read_to_string(out_dir.path().join("bindings/__init__.py"))?;
+
+        assert!(!generated.contains("def x() -> None:"));
+
+        Ok(())
+    }
+
+    #[test]
+    fn unstable_bindings_generated_with_feature_flag() -> Result<()> {
+        // Given a WIT file with gated features
+        let wit = gated_x_wit_file()?;
+        let out_dir = tempfile::tempdir()?;
+
+        // When generating the bindings for this WIT world
+        let common = Common {
+            wit_path: Some(wit.path().into()),
+            world: None,
+            quiet: false,
+            features: vec!["x".to_owned()],
+            all_features: false,
+        };
+        let bindings = Bindings {
+            output_dir: out_dir.path().into(),
+            world_module: None,
+        };
+        generate_bindings(common, bindings)?;
+
+        // Then the gated feature doesn't appear
+        let generated = fs::read_to_string(out_dir.path().join("bindings/__init__.py"))?;
+
+        assert!(generated.contains("def x() -> None:"));
+
+        Ok(())
+    }
+
+    #[test]
+    fn unstable_bindings_generated_for_all_features() -> Result<()> {
+        // Given a WIT file with gated features
+        let wit = gated_x_wit_file()?;
+        let out_dir = tempfile::tempdir()?;
+
+        // When generating the bindings for this WIT world
+        let common = Common {
+            wit_path: Some(wit.path().into()),
+            world: None,
+            quiet: false,
+            features: vec![],
+            all_features: true,
+        };
+        let bindings = Bindings {
+            output_dir: out_dir.path().into(),
+            world_module: None,
+        };
+        generate_bindings(common, bindings)?;
+
+        // Then the gated feature doesn't appear
+        let generated = fs::read_to_string(out_dir.path().join("bindings/__init__.py"))?;
+
+        assert!(generated.contains("def x() -> None:"));
+
+        Ok(())
+    }
+
+    #[test]
+    fn unstable_features_used_in_componentize() -> Result<()> {
+        // Given bindings to a WIT file with gated features and a Python file that uses them
+        let wit = gated_x_wit_file()?;
+        let out_dir = tempfile::tempdir()?;
+        let common = Common {
+            wit_path: Some(wit.path().into()),
+            world: None,
+            quiet: false,
+            features: vec!["x".to_owned()],
+            all_features: false,
+        };
+        let bindings = Bindings {
+            output_dir: out_dir.path().into(),
+            world_module: None,
+        };
+        generate_bindings(common.clone(), bindings)?;
+        fs::write(
+            out_dir.path().join("app.py"),
+            r#"
+import bindings
+from bindings import x
+
+class Bindings(bindings.Bindings):
+    def y(self) -> None:
+        x()
+"#,
+        )?;
+
+        // Building the component succeeds
+        let componentize_opts = Componentize {
+            app_name: "app".to_owned(),
+            python_path: vec![out_dir.path().to_string_lossy().into()],
+            module_worlds: vec![],
+            output: out_dir.path().join("app.wasm"),
+            stub_wasi: false,
+        };
+        componentize(common, componentize_opts)
+    }
+}

src/lib.rs

@@ -170,13 +170,15 @@ impl Invoker for MyInvoker {
 pub fn generate_bindings(
     wit_path: &Path,
     world: Option<&str>,
+    features: &[String],
+    all_features: bool,
     world_module: Option<&str>,
     output_dir: &Path,
 ) -> Result<()> {
     // TODO: Split out and reuse the code responsible for finding and using componentize-py.toml files in the
     // `componentize` function below, since that can affect the bindings we should be generating.
 
-    let (resolve, world) = parse_wit(wit_path, world)?;
+    let (resolve, world) = parse_wit(wit_path, world, features, all_features)?;
     let summary = Summary::try_new(&resolve, &iter::once(world).collect())?;
     let world_name = resolve.worlds[world].name.to_snake_case().escape();
     let world_module = world_module.unwrap_or(&world_name);
@@ -197,6 +199,8 @@ pub fn generate_bindings(
 pub async fn componentize(
     wit_path: Option<&Path>,
     world: Option<&str>,
+    features: &[String],
+    all_features: bool,
     python_path: &[&str],
     module_worlds: &[(&str, &str)],
     app_name: &str,
@@ -219,7 +223,7 @@ pub async fn componentize(
     // Next, iterate over all the WIT directories, merging them into a single `Resolve`, and matching Python
     // packages to `WorldId`s.
     let (mut resolve, mut main_world) = if let Some(path) = wit_path {
-        let (resolve, world) = parse_wit(path, world)?;
+        let (resolve, world) = parse_wit(path, world, features, all_features)?;
         (Some(resolve), Some(world))
     } else {
         (None, None)
@@ -230,7 +234,7 @@ pub async fn componentize(
         .map(|(module, (config, world))| {
             Ok((module, match (world, config.config.wit_directory.as_deref()) {
                 (_, Some(wit_path)) => {
-                    let (my_resolve, mut world) = parse_wit(&config.path.join(wit_path), *world)?;
+                    let (my_resolve, mut world) = parse_wit(&config.path.join(wit_path), *world, features, all_features)?;
 
                     if let Some(resolve) = &mut resolve {
                         let remap = resolve.merge(my_resolve)?;
@@ -254,10 +258,11 @@ pub async fn componentize(
     } else {
         // If no WIT directory was provided as a parameter and none were referenced by Python packages, use ./wit
         // by default.
-        let (my_resolve, world) = parse_wit(Path::new("wit"), world).context(
-            "no WIT files found; please specify the directory or file \
+        let (my_resolve, world) = parse_wit(Path::new("wit"), world, features, all_features)
+            .context(
+                "no WIT files found; please specify the directory or file \
              containing the WIT world you wish to target",
-        )?;
+            )?;
         main_world = Some(world);
         my_resolve
     };
@@ -558,8 +563,25 @@ pub async fn componentize(
     Ok(())
 }
 
-fn parse_wit(path: &Path, world: Option<&str>) -> Result<(Resolve, WorldId)> {
-    let mut resolve = Resolve::default();
+fn parse_wit(
+    path: &Path,
+    world: Option<&str>,
+    features: &[String],
+    all_features: bool,
+) -> Result<(Resolve, WorldId)> {
+    let mut resolve = Resolve {
+        all_features,
+        ..Default::default()
+    };
+    for features in features {
+        for feature in features
+            .split(',')
+            .flat_map(|s| s.split_whitespace())
+            .filter(|f| !f.is_empty())
+        {
+            resolve.features.insert(feature.to_string());
+        }
+    }
     let pkg = if path.is_dir() {
         resolve.push_dir(path)?.0
     } else {

src/python.rs

@@ -7,10 +7,12 @@ use {
 #[allow(clippy::too_many_arguments)]
 #[pyo3::pyfunction]
 #[pyo3(name = "componentize")]
-#[pyo3(signature = (wit_path, world, python_path, module_worlds, app_name, output_path, stub_wasi))]
+#[pyo3(signature = (wit_path, world, features, all_features, python_path, module_worlds, app_name, output_path, stub_wasi))]
 fn python_componentize(
     wit_path: Option<PathBuf>,
     world: Option<&str>,
+    features: Vec<String>,
+    all_features: bool,
     python_path: Vec<&str>,
     module_worlds: Vec<(&str, &str)>,
     app_name: &str,
@@ -21,6 +23,8 @@ fn python_componentize(
         Runtime::new()?.block_on(crate::componentize(
             wit_path.as_deref(),
             world,
+            &features,
+            all_features,
             &python_path,
             &module_worlds,
             app_name,
@@ -34,15 +38,24 @@ fn python_componentize(
 
 #[pyo3::pyfunction]
 #[pyo3(name = "generate_bindings")]
-#[pyo3(signature = (wit_path, world, world_module, output_dir))]
+#[pyo3(signature = (wit_path, world, features, all_features, world_module, output_dir))]
 fn python_generate_bindings(
     wit_path: PathBuf,
     world: Option<&str>,
+    features: Vec<String>,
+    all_features: bool,
     world_module: Option<&str>,
     output_dir: PathBuf,
 ) -> PyResult<()> {
-    crate::generate_bindings(&wit_path, world, world_module, &output_dir)
-        .map_err(|e| PyAssertionError::new_err(format!("{e:?}")))
+    crate::generate_bindings(
+        &wit_path,
+        world,
+        &features,
+        all_features,
+        world_module,
+        &output_dir,
+    )
+    .map_err(|e| PyAssertionError::new_err(format!("{e:?}")))
 }
 
 #[pyo3::pyfunction]

src/test.rs

@@ -63,6 +63,8 @@ async fn make_component(
     crate::componentize(
         Some(&tempdir.path().join("app.wit")),
         None,
+        &[],
+        false,
         &python_path
             .iter()
             .copied()