Skip to content

Commit 6469f8f

Browse files
committed
Support #[cfg] in with
1 parent 5bf13e8 commit 6469f8f

File tree

6 files changed

+156
-47
lines changed

6 files changed

+156
-47
lines changed

crates/derive-impl/src/lib.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ use quote::ToTokens;
2626
use syn::{DeriveInput, Item};
2727
use syn_ext::types::PunctuatedNestedMeta;
2828

29+
pub use pymodule::PyModuleArgs;
30+
2931
pub use compile_bytecode::Compiler;
3032

3133
fn result_to_tokens(result: Result<TokenStream, impl Into<Diagnostic>>) -> TokenStream {
@@ -54,7 +56,7 @@ pub fn pyexception(attr: PunctuatedNestedMeta, item: Item) -> TokenStream {
5456
}
5557
}
5658

57-
pub fn pymodule(attr: PunctuatedNestedMeta, item: Item) -> TokenStream {
59+
pub fn pymodule(attr: PyModuleArgs, item: Item) -> TokenStream {
5860
result_to_tokens(pymodule::impl_pymodule(attr, item))
5961
}
6062

crates/derive-impl/src/pymodule.rs

Lines changed: 144 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,97 @@ use crate::util::{
77
};
88
use core::str::FromStr;
99
use proc_macro2::{Delimiter, Group, TokenStream, TokenTree};
10-
use quote::{ToTokens, quote, quote_spanned};
10+
use quote::{ToTokens, format_ident, quote, quote_spanned};
1111
use rustpython_doc::DB;
1212
use std::collections::HashSet;
1313
use syn::{Attribute, Ident, Item, Result, parse_quote, spanned::Spanned};
1414
use syn_ext::ext::*;
15-
use syn_ext::types::PunctuatedNestedMeta;
15+
use syn_ext::types::NestedMeta;
16+
17+
/// A `with(...)` item that may be gated by `#[cfg(...)]` attributes.
18+
pub struct WithItem {
19+
pub cfg_attrs: Vec<Attribute>,
20+
pub path: syn::Path,
21+
}
22+
23+
impl syn::parse::Parse for WithItem {
24+
fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
25+
let cfg_attrs = Attribute::parse_outer(input)?;
26+
for attr in &cfg_attrs {
27+
if !attr.path().is_ident("cfg") {
28+
return Err(syn::Error::new_spanned(
29+
attr,
30+
"only #[cfg(...)] is supported in with()",
31+
));
32+
}
33+
}
34+
let path = input.parse()?;
35+
Ok(WithItem { cfg_attrs, path })
36+
}
37+
}
38+
39+
/// Parsed arguments for `#[pymodule(...)]`, supporting `#[cfg]` inside `with(...)`.
40+
pub struct PyModuleArgs {
41+
pub metas: Vec<NestedMeta>,
42+
pub with_items: Vec<WithItem>,
43+
}
44+
45+
impl syn::parse::Parse for PyModuleArgs {
46+
fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
47+
let mut metas = Vec::new();
48+
let mut with_items = Vec::new();
49+
50+
while !input.is_empty() {
51+
// Detect `with(...)` — an ident "with" followed by a paren group
52+
if input.peek(Ident) && input.peek2(syn::token::Paren) {
53+
let fork = input.fork();
54+
let ident: Ident = fork.parse()?;
55+
if ident == "with" {
56+
// Advance past "with"
57+
let _: Ident = input.parse()?;
58+
let content;
59+
syn::parenthesized!(content in input);
60+
let items =
61+
syn::punctuated::Punctuated::<WithItem, syn::Token![,]>::parse_terminated(
62+
&content,
63+
)?;
64+
with_items = items.into_iter().collect();
65+
if !input.is_empty() {
66+
input.parse::<syn::Token![,]>()?;
67+
}
68+
continue;
69+
}
70+
}
71+
metas.push(input.parse::<NestedMeta>()?);
72+
if input.is_empty() {
73+
break;
74+
}
75+
input.parse::<syn::Token![,]>()?;
76+
}
77+
78+
Ok(PyModuleArgs { metas, with_items })
79+
}
80+
}
81+
82+
/// Generate `#[cfg(not(...))]` attributes that negate the given `#[cfg(...)]` attributes.
83+
fn negate_cfg_attrs(cfg_attrs: &[Attribute]) -> Vec<Attribute> {
84+
if cfg_attrs.is_empty() {
85+
return vec![];
86+
}
87+
let preds: Vec<_> = cfg_attrs
88+
.iter()
89+
.map(|attr| match &attr.meta {
90+
syn::Meta::List(list) => list.tokens.clone(),
91+
_ => unreachable!("only #[cfg(...)] should be here"),
92+
})
93+
.collect();
94+
if preds.len() == 1 {
95+
let pred = &preds[0];
96+
vec![parse_quote!(#[cfg(not(#pred))])]
97+
} else {
98+
vec![parse_quote!(#[cfg(not(all(#(#preds),*)))])]
99+
}
100+
}
16101

17102
#[derive(Clone, Copy, Eq, PartialEq)]
18103
enum AttrName {
@@ -62,14 +147,15 @@ struct ModuleContext {
62147
errors: Vec<syn::Error>,
63148
}
64149

65-
pub fn impl_pymodule(attr: PunctuatedNestedMeta, module_item: Item) -> Result<TokenStream> {
150+
pub fn impl_pymodule(args: PyModuleArgs, module_item: Item) -> Result<TokenStream> {
151+
let PyModuleArgs { metas, with_items } = args;
66152
let (doc, mut module_item) = match module_item {
67153
Item::Mod(m) => (m.attrs.doc(), m),
68154
other => bail_span!(other, "#[pymodule] can only be on a full module"),
69155
};
70156
let fake_ident = Ident::new("pymodule", module_item.span());
71157
let module_meta =
72-
ModuleItemMeta::from_nested(module_item.ident.clone(), fake_ident, attr.into_iter())?;
158+
ModuleItemMeta::from_nested(module_item.ident.clone(), fake_ident, metas.into_iter())?;
73159

74160
// generation resources
75161
let mut context = ModuleContext {
@@ -119,7 +205,6 @@ pub fn impl_pymodule(attr: PunctuatedNestedMeta, module_item: Item) -> Result<To
119205
quote!(None)
120206
};
121207
let is_submodule = module_meta.sub()?;
122-
let withs = module_meta.with()?;
123208
if !is_submodule {
124209
items.extend([
125210
parse_quote! {
@@ -154,16 +239,66 @@ pub fn impl_pymodule(attr: PunctuatedNestedMeta, module_item: Item) -> Result<To
154239
}
155240
});
156241
}
157-
let method_defs = if withs.is_empty() {
242+
// Split with_items into unconditional and cfg-gated groups
243+
let (uncond_withs, cond_withs): (Vec<_>, Vec<_>) =
244+
with_items.iter().partition(|w| w.cfg_attrs.is_empty());
245+
let uncond_paths: Vec<_> = uncond_withs.iter().map(|w| &w.path).collect();
246+
247+
let method_defs = if with_items.is_empty() {
158248
quote!(#function_items)
159249
} else {
250+
// For cfg-gated with items, generate conditional const declarations
251+
// so the total array size adapts to the cfg at compile time
252+
let cond_const_names: Vec<_> = cond_withs
253+
.iter()
254+
.enumerate()
255+
.map(|(i, _)| format_ident!("__WITH_METHODS_{}", i))
256+
.collect();
257+
let cond_const_decls: Vec<_> = cond_withs
258+
.iter()
259+
.zip(&cond_const_names)
260+
.map(|(w, name)| {
261+
let cfg_attrs = &w.cfg_attrs;
262+
let neg_attrs = negate_cfg_attrs(&w.cfg_attrs);
263+
let path = &w.path;
264+
quote! {
265+
#(#cfg_attrs)*
266+
const #name: &'static [::rustpython_vm::function::PyMethodDef] = super::#path::METHOD_DEFS;
267+
#(#neg_attrs)*
268+
const #name: &'static [::rustpython_vm::function::PyMethodDef] = &[];
269+
}
270+
})
271+
.collect();
272+
160273
quote!({
161274
const OWN_METHODS: &'static [::rustpython_vm::function::PyMethodDef] = &#function_items;
275+
#(#cond_const_decls)*
162276
rustpython_vm::function::PyMethodDef::__const_concat_arrays::<
163-
{ OWN_METHODS.len() #(+ super::#withs::METHOD_DEFS.len())* },
164-
>(&[#(super::#withs::METHOD_DEFS,)* OWN_METHODS])
277+
{ OWN_METHODS.len()
278+
#(+ super::#uncond_paths::METHOD_DEFS.len())*
279+
#(+ #cond_const_names.len())*
280+
},
281+
>(&[
282+
#(super::#uncond_paths::METHOD_DEFS,)*
283+
#(#cond_const_names,)*
284+
OWN_METHODS
285+
])
165286
})
166287
};
288+
289+
// Generate __init_attributes calls, wrapping cfg-gated items
290+
let init_with_calls: Vec<_> = with_items
291+
.iter()
292+
.map(|w| {
293+
let cfg_attrs = &w.cfg_attrs;
294+
let path = &w.path;
295+
quote! {
296+
#(#cfg_attrs)*
297+
super::#path::__init_attributes(vm, module);
298+
}
299+
})
300+
.collect();
301+
167302
items.extend([
168303
parse_quote! {
169304
::rustpython_vm::common::static_cell! {
@@ -178,9 +313,7 @@ pub fn impl_pymodule(attr: PunctuatedNestedMeta, module_item: Item) -> Result<To
178313
vm: &::rustpython_vm::VirtualMachine,
179314
module: &::rustpython_vm::Py<::rustpython_vm::builtins::PyModule>,
180315
) {
181-
#(
182-
super::#withs::__init_attributes(vm, module);
183-
)*
316+
#(#init_with_calls)*
184317
let ctx = &vm.ctx;
185318
#attribute_items
186319
}

crates/derive-impl/src/util.rs

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -315,7 +315,7 @@ impl ItemMeta for SimpleItemMeta {
315315
pub(crate) struct ModuleItemMeta(pub ItemMetaInner);
316316

317317
impl ItemMeta for ModuleItemMeta {
318-
const ALLOWED_NAMES: &'static [&'static str] = &["name", "with", "sub"];
318+
const ALLOWED_NAMES: &'static [&'static str] = &["name", "sub"];
319319

320320
fn from_inner(inner: ItemMetaInner) -> Self {
321321
Self(inner)
@@ -330,20 +330,6 @@ impl ModuleItemMeta {
330330
pub fn sub(&self) -> Result<bool> {
331331
self.inner()._bool("sub")
332332
}
333-
334-
pub fn with(&self) -> Result<Vec<&syn::Path>> {
335-
let mut withs = Vec::new();
336-
let Some(nested) = self.inner()._optional_list("with")? else {
337-
return Ok(withs);
338-
};
339-
for meta in nested {
340-
let NestedMeta::Meta(Meta::Path(path)) = meta else {
341-
bail_span!(meta, "#[pymodule(with(...))] arguments should be paths")
342-
};
343-
withs.push(path);
344-
}
345-
Ok(withs)
346-
}
347333
}
348334

349335
pub(crate) struct AttrItemMeta(pub ItemMetaInner);

crates/derive/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@ pub fn pyexception(attr: TokenStream, item: TokenStream) -> TokenStream {
209209
/// - `name`: the name of the function in Python, by default it is the same as the associated Rust function.
210210
#[proc_macro_attribute]
211211
pub fn pymodule(attr: TokenStream, item: TokenStream) -> TokenStream {
212-
let attr = parse_macro_input!(attr with Punctuated::parse_terminated);
212+
let attr = parse_macro_input!(attr as derive_impl::PyModuleArgs);
213213
let item = parse_macro_input!(item);
214214
derive_impl::pymodule(attr, item).into()
215215
}

crates/stdlib/src/openssl.rs

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,12 @@ fn probe() -> &'static ProbeResult {
4343
}
4444

4545
#[allow(non_upper_case_globals)]
46-
#[pymodule(with(cert::ssl_cert, ssl_error::ssl_error, ossl101, ossl111, windows))]
46+
#[pymodule(with(
47+
cert::ssl_cert,
48+
ssl_error::ssl_error,
49+
#[cfg(ossl101)] ossl101,
50+
#[cfg(ossl111)] ossl111,
51+
#[cfg(windows)] windows))]
4752
mod _ssl {
4853
use super::{bio, probe};
4954

@@ -4070,18 +4075,6 @@ mod _ssl {
40704075
}
40714076
}
40724077

4073-
#[cfg(not(ossl101))]
4074-
#[pymodule(sub)]
4075-
mod ossl101 {}
4076-
4077-
#[cfg(not(ossl111))]
4078-
#[pymodule(sub)]
4079-
mod ossl111 {}
4080-
4081-
#[cfg(not(windows))]
4082-
#[pymodule(sub)]
4083-
mod windows {}
4084-
40854078
#[allow(non_upper_case_globals)]
40864079
#[cfg(ossl101)]
40874080
#[pymodule(sub)]

crates/vm/src/stdlib/time.rs

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ unsafe extern "C" {
2323
fn c_tzset();
2424
}
2525

26-
#[pymodule(name = "time", with(platform))]
26+
#[pymodule(name = "time", with(#[cfg(any(unix, windows))] platform))]
2727
mod decl {
2828
use crate::{
2929
AsObject, Py, PyObjectRef, PyResult, VirtualMachine,
@@ -986,8 +986,3 @@ mod platform {
986986
Ok(Duration::from_nanos((k_time + u_time) * 100))
987987
}
988988
}
989-
990-
// mostly for wasm32
991-
#[cfg(not(any(unix, windows)))]
992-
#[pymodule(sub)]
993-
mod platform {}

0 commit comments

Comments
 (0)