|
10 | 10 | use allocative::Allocative; |
11 | 11 | use buck2_core::package::PackageLabel; |
12 | 12 | use buck2_interpreter::types::configured_providers_label::StarlarkProvidersLabel; |
| 13 | +use buck2_node::attrs::coerced_attr::CoercedConcat; |
13 | 14 | use buck2_node::attrs::coerced_attr::CoercedSelector; |
14 | 15 | use buck2_node::attrs::coerced_attr::CoercedSelectorKeyRef; |
15 | 16 | use buck2_node::attrs::display::AttrDisplayWithContext; |
@@ -216,3 +217,114 @@ fn select_dict_methods(builder: &mut MethodsBuilder) { |
216 | 217 | Ok(this.get(key, heap)?) |
217 | 218 | } |
218 | 219 | } |
| 220 | + |
| 221 | +#[derive(ProvidesStaticType, Derivative, Trace, Allocative, Clone, Debug)] |
| 222 | +pub(crate) struct StarlarkSelectConcat { |
| 223 | + concat: CoercedConcat, |
| 224 | + pkg: PackageLabel, |
| 225 | +} |
| 226 | + |
| 227 | +impl StarlarkSelectConcat { |
| 228 | + pub(crate) fn new(concat: CoercedConcat, pkg: PackageLabel) -> Self { |
| 229 | + Self { concat, pkg } |
| 230 | + } |
| 231 | +} |
| 232 | + |
| 233 | +impl Serialize for StarlarkSelectConcat { |
| 234 | + fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> |
| 235 | + where |
| 236 | + S: Serializer, |
| 237 | + { |
| 238 | + let ctx = AttrFmtContext { |
| 239 | + package: Some(self.pkg), |
| 240 | + options: Default::default(), |
| 241 | + }; |
| 242 | + self.concat.serialize_with_ctx(&ctx, serializer) |
| 243 | + } |
| 244 | +} |
| 245 | + |
| 246 | +impl Display for StarlarkSelectConcat { |
| 247 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 248 | + let ctx = AttrFmtContext { |
| 249 | + package: Some(self.pkg), |
| 250 | + options: Default::default(), |
| 251 | + }; |
| 252 | + self.concat.fmt(&ctx, f) |
| 253 | + } |
| 254 | +} |
| 255 | + |
| 256 | +starlark_simple_value!(StarlarkSelectConcat); |
| 257 | + |
| 258 | +#[starlark_value(type = "bxl.SelectConcat")] |
| 259 | +impl<'v> StarlarkValue<'v> for StarlarkSelectConcat { |
| 260 | + fn get_methods() -> Option<&'static Methods> { |
| 261 | + static RES: MethodsStatic = MethodsStatic::new(); |
| 262 | + RES.methods(select_concat_methods) |
| 263 | + } |
| 264 | + |
| 265 | + fn length(&self) -> starlark::Result<i32> { |
| 266 | + Ok(self.concat.0.len() as i32) |
| 267 | + } |
| 268 | +} |
| 269 | + |
| 270 | +/// In bxl, `Select = bxl.SelectDict | bxl.SelectConcat`. `bxl.SelectConcat` is a list-like object that represents a select. |
| 271 | +/// One example of this type is: |
| 272 | +/// ``` |
| 273 | +/// ["--flags"] + select({ |
| 274 | +// "root//constraints:a": ["--foo"], |
| 275 | +// "root//constraints:b": ["--bar"], |
| 276 | +// "DEFAULT": ["baz"] |
| 277 | +// }) |
| 278 | +/// ``` |
| 279 | +/// You can: |
| 280 | +/// * Iterate over the values of this object (e.g. `for item in select_concat.select_iter():`) |
| 281 | +/// * Get the length (e.g. `len(select_concat)`) |
| 282 | +/// * Check its type using `isinstance(select_concat, bxl.SelectConcat)`. |
| 283 | +/// |
| 284 | +/// Simple usage: |
| 285 | +/// ```python |
| 286 | +/// def _impl_select_concat(ctx): |
| 287 | +/// node = ctx.lazy.unconfigured_target_node("root//:select_concat").resolve() |
| 288 | +/// attr = node.get_attr("select_attr") |
| 289 | +/// for value in attr: |
| 290 | +/// if isinstance(value, bxl.SelectDict): |
| 291 | +/// for key, value in value.items(): |
| 292 | +/// ctx.output.print(f"{key} -> {value}") |
| 293 | +/// else: |
| 294 | +/// ctx.output.print(value) |
| 295 | +/// ctx.output.print(attr[0]) |
| 296 | +/// ``` |
| 297 | +#[starlark_module] |
| 298 | +fn select_concat_methods(builder: &mut MethodsBuilder) { |
| 299 | + /// Return the values of the SelectConcat. |
| 300 | + /// |
| 301 | + /// Sample usage: |
| 302 | + /// ```python |
| 303 | + /// def _impl_select_concat(ctx): |
| 304 | + /// node = ctx.lazy.unconfigured_target_node("root//:select_concat").resolve() |
| 305 | + /// attr = node.get_attr("select_attr") |
| 306 | + /// for value in attr.select_iter(): |
| 307 | + /// ctx.output.print(value) |
| 308 | + /// ``` |
| 309 | + fn select_iter<'v>( |
| 310 | + this: &'v StarlarkSelectConcat, |
| 311 | + heap: &'v Heap, |
| 312 | + ) -> anyhow::Result<Vec<Value<'v>>> { |
| 313 | + let list = this |
| 314 | + .concat |
| 315 | + .iter() |
| 316 | + .map(|a| a.to_value(this.pkg.dupe(), heap)) |
| 317 | + .collect::<buck2_error::Result<_>>()?; |
| 318 | + Ok(list) |
| 319 | + } |
| 320 | + |
| 321 | + /// Returns the length of a SelectConcat, defined as the number of items being concatenated |
| 322 | + /// at the select level (not the total number of elements across all lists). |
| 323 | + /// |
| 324 | + /// For example, `[1, 2] + select({"DEFAULT": [3, 4]}` returns 2 instead of 4. |
| 325 | + /// Note: You can use `len()` to get the length too. |
| 326 | + #[starlark(attribute)] |
| 327 | + fn length(this: &StarlarkSelectConcat) -> starlark::Result<i32> { |
| 328 | + Ok(this.concat.0.len() as i32) |
| 329 | + } |
| 330 | +} |
0 commit comments