feat: file.symbols and file.get_symbol (#53)

* Use supertypes

* Fix path to wheels

* Refactor

* Revert change

* Add symbols

* Generate get_symbol and symbols

* Fix bugs

* Don't add symbol enums to the module

* Remove unused import

* update snapshots

Author: bagel897

Cargo.lock

@@ -141,6 +141,7 @@ indextree = { version = "4.7.3", features = ["std"], default-features = false }
 thiserror = "2.0.11"
 indexmap = "2"
 smallvec = "1.11.0"
+pluralizer = "0.5.0"
 
 
 [profile.dev]

Cargo.toml

@@ -19,6 +19,8 @@ syn = { workspace = true }
 anyhow = { workspace = true }
 proc-macro2 = { workspace = true }
 codegen-sdk-cst = { workspace = true }
+pluralizer = "0.5.0"
+convert_case = { workspace = true }
 [dev-dependencies]
 test-log = { workspace = true }
 rstest = { workspace = true }

codegen-bindings-generator/Cargo.toml

@@ -1,11 +1,85 @@
 use codegen_sdk_ast_generator::{HasQuery, Symbol};
 use codegen_sdk_common::Language;
 use codegen_sdk_cst::CSTDatabase;
+use convert_case::{Case, Casing};
+use pluralizer::pluralize;
 use proc_macro2::Span;
 use quote::{format_ident, quote};
 use syn::{parse_quote, parse_quote_spanned};
 
 use super::{cst::generate_cst, helpers};
+fn get_category(group: &str) -> Vec<syn::Stmt> {
+    let category = format_ident!("{}", pluralize(group, 2, false));
+    parse_quote! {
+        let file = self.file(py)?;
+        let db = self.codebase.get(py).db();
+        let category = file.#category(db);
+    }
+}
+fn filter_symbols(
+    nodes: &Vec<&codegen_sdk_ast_generator::Symbol>,
+    group: &str,
+) -> Vec<codegen_sdk_ast_generator::Symbol> {
+    nodes
+        .iter()
+        .filter(|node| node.category == pluralize(group, 2, false))
+        .cloned()
+        .cloned()
+        .collect()
+}
+fn get_symbol_method_name(group: &str) -> syn::Ident {
+    let symbol_name = codegen_sdk_ast_generator::get_symbol_name(group);
+    syn::Ident::new(
+        &format!("get_{}", symbol_name.to_string().to_case(Case::Snake)),
+        Span::call_site(),
+    )
+}
+// Generates the get_symbol method
+fn get_symbols_method(group: &str) -> Vec<syn::Stmt> {
+    let mut output: Vec<syn::Stmt> = Vec::new();
+    let symbol_name = codegen_sdk_ast_generator::get_symbol_name(group);
+    let category = get_category(group);
+    let symbols_method = codegen_sdk_ast_generator::get_symbols_method(&symbol_name);
+    let get_symbols_method = get_symbol_method_name(group);
+    output.extend::<Vec<syn::Stmt>>(parse_quote! {
+        #[pyo3(signature = (name,optional=false))]
+        pub fn #get_symbols_method(&self, py: Python<'_>, name: String, optional: bool) -> PyResult<Option<#symbol_name>> {
+            #(#category)*
+            let subcategory = category.#symbols_method(db);
+            let res = subcategory.get(&name);
+            if let Some(nodes) = res {
+                if nodes.len() == 1 {
+                    Ok(Some(#symbol_name::new(py.clone(), nodes[0].fully_qualified_name(db), 0, &nodes[0], self.codebase.clone())))
+                } else {
+                    Err(pyo3::exceptions::PyValueError::new_err(format!("Ambiguous symbol {} found {} possible matches", name, nodes.len())))
+                }
+            } else {
+                if optional {
+                    Ok(None)
+                } else {
+                    Err(pyo3::exceptions::PyValueError::new_err(format!("No symbol {} found", name)))
+                }
+            }
+    }
+    });
+    output
+}
+fn symbols_method(group: &str) -> Vec<syn::Stmt> {
+    let mut output: Vec<syn::Stmt> = Vec::new();
+    let symbol_name = codegen_sdk_ast_generator::get_symbol_name(group);
+    let category = get_category(group);
+    let symbols_method = codegen_sdk_ast_generator::get_symbols_method(&symbol_name);
+    output.extend::<Vec<syn::Stmt>>(parse_quote! {
+        #[getter]
+        pub fn #symbols_method(&self, py: Python<'_>) -> PyResult<Vec<#symbol_name>> {
+            #(#category)*
+            let subcategory = category.#symbols_method(db);
+            let nodes = subcategory.values().map(|values| values.into_iter().enumerate().map(|(idx, node)| #symbol_name::new(py.clone(), node.fully_qualified_name(db), idx, node, self.codebase.clone()))).flatten().collect();
+            Ok(nodes)
+    }
+    });
+    output
+}
 fn generate_file_struct(
     language: &Language,
     symbols: Vec<&Symbol>,
@@ -47,6 +121,13 @@ fn generate_file_struct(
         .filter(|symbol| symbol.category != symbol.subcategory)
         .map(|symbol| vec![symbol.py_file_getter(), symbol.py_file_get()])
         .flatten();
+    let mut symbols_methods = Vec::new();
+    for group in codegen_sdk_ast_generator::GROUPS {
+        if filter_symbols(&symbols, group).len() > 0 {
+            symbols_methods.extend(symbols_method(group));
+            symbols_methods.extend(get_symbols_method(group));
+        }
+    }
     output.push(parse_quote! {
         #[pymethods]
         impl #struct_name {
@@ -70,6 +151,7 @@ fn generate_file_struct(
                 Ok(self.content(py)?.to_string())
             }
             #(#methods)*
+            #(#symbols_methods)*
         }
     });
     Ok(output)
@@ -158,6 +240,54 @@ fn generate_symbol_struct(
     });
     Ok(output)
 }
+// Generate an enum for all possible symbols. (IE Symbol => Class, Function, etc)
+fn generate_symbol_enum(
+    language: &Language,
+    symbols: Vec<&Symbol>,
+    group: &str,
+) -> anyhow::Result<Vec<syn::Stmt>> {
+    let symbols: Vec<_> = filter_symbols(&symbols, group)
+        .iter()
+        .map(|symbol| format_ident!("{}", symbol.name))
+        .collect();
+    if symbols.len() == 0 {
+        return Ok(Vec::new());
+    }
+    let symbol_name = codegen_sdk_ast_generator::get_symbol_name(group);
+    let span = Span::call_site();
+    let mut output = Vec::new();
+    let enum_name = format_ident!("{}", symbol_name);
+    let package_name = syn::Ident::new(&language.package_name(), span);
+    output.push(parse_quote_spanned! {
+        span =>
+        #[derive(IntoPyObject)]
+        pub enum #enum_name {
+            #(#symbols(#symbols),)*
+        }
+    });
+    let original_name = quote! { codegen_sdk_analyzer::#package_name::ast::#symbol_name };
+    let matchers: Vec<syn::Arm> = symbols
+        .iter()
+        .map(|symbol| {
+            let symbol_name = format_ident!("{}", symbol);
+            parse_quote_spanned! {
+                span =>
+                #original_name::#symbol_name(_) => Self::#symbol_name(#symbol_name::new(id, idx, codebase_arc)),
+            }
+        })
+        .collect();
+    output.push(parse_quote_spanned! {
+        span =>
+        impl #enum_name {
+            pub fn new(py: Python<'_>, id: codegen_sdk_resolution::FullyQualifiedName, idx: usize, node: &#original_name<'_>, codebase_arc: Arc<GILProtected<codegen_sdk_analyzer::Codebase>>) -> Self {
+                match node {
+                    #(#matchers)*
+                }
+            }
+        }
+    });
+    Ok(output)
+}
 fn generate_module(
     language: &Language,
     symbols: Vec<syn::Ident>,
@@ -167,6 +297,7 @@ fn generate_module(
     let register_name = format_ident!("register_{}", language_name);
     let struct_name = format_ident!("{}", language.file_struct_name());
     let module_name = format!("codegen_sdk_pink.{}", language_name);
+
     output.push(parse_quote! {
         pub fn #register_name(py: Python<'_>, parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
             let child_module = PyModule::new(parent_module.py(), #language_name)?;
@@ -194,6 +325,10 @@ pub(crate) fn generate_bindings(language: &Language) -> anyhow::Result<Vec<syn::
     let cst = generate_cst(language, &state)?;
     output.extend(cst);
     let mut symbol_idents = Vec::new();
+    for group in codegen_sdk_ast_generator::GROUPS {
+        let symbol_enum = generate_symbol_enum(language, symbols.values().collect(), group)?;
+        output.extend(symbol_enum);
+    }
     for (_, symbol) in symbols {
         let symbol_struct = generate_symbol_struct(language, &symbol)?;
         output.extend(symbol_struct);