Subenum generation

Author: bagel897

codegen-bindings-generator/src/python/cst.rs

@@ -1,6 +1,6 @@
-use codegen_sdk_common::Language;
+use codegen_sdk_common::{Language, naming::normalize_type_name};
 use proc_macro2::Span;
-use quote::format_ident;
+use quote::{format_ident, quote};
 use syn::parse_quote;
 
 use super::helpers;
@@ -19,7 +19,7 @@ fn generate_cst_struct(
             codebase: Arc<GILProtected<codegen_sdk_analyzer::Codebase>>,
         }
     });
-    let file_getter = helpers::get_file(language);
+    let file_getter = helpers::get_file(language, quote! { self.id }, quote! { self.codebase });
     output.push(parse_quote! {
         impl #struct_name {
             pub fn new(id: codegen_sdk_common::CSTNodeTreeId, codebase: Arc<GILProtected<codegen_sdk_analyzer::Codebase>>) -> Self {
@@ -87,6 +87,65 @@ fn generate_cst_struct(
     });
     Ok(output)
 }
+fn generate_cst_subenum(
+    language: &Language,
+    state: &codegen_sdk_cst_generator::State,
+    name: &str,
+) -> anyhow::Result<Vec<syn::Stmt>> {
+    let mut output = Vec::new();
+    let struct_name = format_ident!("{}", normalize_type_name(name, true));
+    let package_name = syn::Ident::new(&language.package_name(), Span::call_site());
+    let module_name = format!("codegen_sdk_pink::{}.cst", language.name());
+    let subenum_names = state
+        .get_subenum_variants(&name, false)
+        .iter()
+        .map(|name| {
+            let name = format_ident!("{}", name.normalize_name());
+            parse_quote! {
+                #name(#name)
+            }
+        })
+        .collect::<Vec<syn::Variant>>();
+    output.push(parse_quote! {
+        #[derive(IntoPyObject)]
+        pub enum #struct_name {
+            #(#subenum_names,)*
+        }
+    });
+    let package_name = syn::Ident::new(&language.package_name(), Span::call_site());
+    let ref_name = syn::Ident::new(
+        &format!("{}Ref", struct_name.to_string()),
+        Span::call_site(),
+    );
+    let matchers: Vec<syn::Arm> = state
+        .get_subenum_variants(&name, false)
+        .iter()
+        .map(|node| {
+            let name = format_ident!("{}", node.normalize_name());
+            parse_quote! {
+                codegen_sdk_analyzer::#package_name::cst::#ref_name::#name(_) => Ok(Self::#name(#name::new(id, codebase_arc.clone()))),
+            }
+        })
+        .collect();
+    let get_file = helpers::get_file(language, quote! { id }, quote! { codebase_arc });
+    output.push(parse_quote! {
+        impl #struct_name {
+            pub fn new(py: Python<'_>, id: codegen_sdk_common::CSTNodeTreeId, codebase_arc: Arc<GILProtected<codegen_sdk_analyzer::Codebase>>) -> PyResult<Self> {
+                #(#get_file)*
+                let node = file.tree(codebase.db()).get(id.id(codebase.db()));
+                if let Some(node) = node {
+                    match node.as_ref().try_into().unwrap() {
+                        #(#matchers)*
+                    }
+                } else {
+                    Err(pyo3::exceptions::PyValueError::new_err("Node not found"))
+                }
+            }
+        }
+    });
+    Ok(output)
+}
+
 fn generate_module(state: &codegen_sdk_cst_generator::State) -> anyhow::Result<Vec<syn::Stmt>> {
     let mut output = Vec::new();
     let node_names = state.get_node_struct_names();
@@ -110,6 +169,10 @@ pub fn generate_cst(
         let cst_struct = generate_cst_struct(language, node)?;
         output.extend(cst_struct);
     }
+    for subenum in &state.subenums {
+        let cst_subenum = generate_cst_subenum(language, state, subenum)?;
+        output.extend(cst_subenum);
+    }
     output.extend(generate_module(state)?);
     Ok(parse_quote! {
         mod cst {

codegen-bindings-generator/src/python/generator.rs

@@ -2,7 +2,7 @@ use codegen_sdk_ast_generator::{HasQuery, Symbol};
 use codegen_sdk_common::Language;
 use codegen_sdk_cst::CSTDatabase;
 use proc_macro2::Span;
-use quote::format_ident;
+use quote::{format_ident, quote};
 use syn::{parse_quote, parse_quote_spanned};
 
 use super::{cst::generate_cst, helpers};
@@ -92,7 +92,7 @@ fn generate_symbol_struct(
             codebase: Arc<GILProtected<codegen_sdk_analyzer::Codebase>>,
         }
     });
-    let file_getter = helpers::get_file(language);
+    let file_getter = helpers::get_file(language, quote! { self.id }, quote! { self.codebase });
     let category = syn::Ident::new(&symbol.category, span);
     let subcategory = syn::Ident::new(&symbol.subcategory, span);
     output.push(parse_quote_spanned! {

codegen-bindings-generator/src/python/helpers.rs

@@ -1,13 +1,13 @@
 use codegen_sdk_common::Language;
-use proc_macro2::Span;
+use proc_macro2::{Span, TokenStream};
 use syn::parse_quote_spanned;
-pub fn get_file(language: &Language) -> Vec<syn::Stmt> {
+pub fn get_file(language: &Language, id: TokenStream, codebase: TokenStream) -> Vec<syn::Stmt> {
     let span = Span::call_site();
     let variant_name = syn::Ident::new(&language.struct_name, span);
     parse_quote_spanned! {
     span =>
-    let codebase = self.codebase.get(py);
-    let path = self.id.file(codebase.db());
+    let codebase = #codebase.get(py);
+    let path = #id.file(codebase.db());
     let file = codebase.get_file_for_id(path);
     let file = match file {
         Some(codegen_sdk_analyzer::ParsedFile::#variant_name(py)) => py,

codegen-sdk-ast-generator/src/query.rs

@@ -619,7 +619,7 @@ impl<'a> Query<'a> {
             );
             matchers.extend_one(matcher);
         } else {
-            let subenum = self.state.get_subenum_variants(&first_node.source());
+            let subenum = self.state.get_subenum_variants(&first_node.source(), false);
             log::info!(
                 "subenum {} with {} variants",
                 first_node.source(),

codegen-sdk-cst-generator/src/generator/state.rs

@@ -153,17 +153,19 @@ impl<'a> State<'a> {
     }
     pub fn get_variants(&self, subenum: &str, include_comment: bool) -> Vec<TypeDefinition> {
         let mut variants = Vec::new();
-        if include_comment {
-            let comment = get_comment_type();
-            variants.push(comment);
-        }
         for node in self.nodes.values() {
             log::debug!("Checking subenum: {} for {}", subenum, node.kind());
             if node.subenums.contains(&subenum.to_string()) {
                 log::debug!("Found variant: {} for {}", node.kind(), subenum);
                 variants.push(node.type_definition());
             }
         }
+        if include_comment {
+            let comment = get_comment_type();
+            if !variants.iter().any(|v| v.type_name == comment.type_name) {
+                variants.push(comment);
+            }
+        }
         variants
     }
     fn get_variant_map(&self, enum_name: &str) -> BTreeMap<u16, TokenStream> {
@@ -290,8 +292,8 @@ impl<'a> State<'a> {
         }
         None
     }
-    pub fn get_subenum_variants(&self, name: &str) -> Vec<&Node<'a>> {
-        let variants = self.get_variants(name, true);
+    pub fn get_subenum_variants(&self, name: &str, include_comment: bool) -> Vec<&Node<'a>> {
+        let variants = self.get_variants(name, include_comment);
         let mut nodes = Vec::new();
         for variant in variants {
             if let Some(node) = self.get_node_for_struct_name(&variant.normalize()) {