@@ -7,12 +7,97 @@ use crate::util::{
77} ;
88use core:: str:: FromStr ;
99use proc_macro2:: { Delimiter , Group , TokenStream , TokenTree } ;
10- use quote:: { ToTokens , quote, quote_spanned} ;
10+ use quote:: { ToTokens , format_ident , quote, quote_spanned} ;
1111use rustpython_doc:: DB ;
1212use std:: collections:: HashSet ;
1313use syn:: { Attribute , Ident , Item , Result , parse_quote, spanned:: Spanned } ;
1414use 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 ) ]
18103enum 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 }
0 commit comments