1#![allow(clippy::style)]
4
5#[cfg(test)]
6mod tests;
7mod unbox;
8
9use std::collections::HashMap;
10
11use convert_case::Casing;
12use inflector::{cases::camelcase::to_camel_case, Inflector};
13use once_cell::sync::Lazy;
14use quote::{format_ident, quote, quote_spanned, ToTokens};
15use regex::Regex;
16use serde::Deserialize;
17use serde_tokenstream::{from_tokenstream, Error};
18use syn::{
19 parse::{Parse, ParseStream},
20 Attribute, Signature, Visibility,
21};
22use unbox::unbox;
23
24#[derive(Deserialize, Debug)]
26struct ArgMetadata {
27 docs: String,
29
30 #[serde(default)]
33 include_in_snippet: bool,
34}
35
36#[derive(Deserialize, Debug)]
37struct StdlibMetadata {
38 name: String,
40
41 #[serde(default)]
43 tags: Vec<String>,
44
45 #[serde(default)]
48 unpublished: bool,
49
50 #[serde(default)]
53 deprecated: bool,
54
55 #[serde(default)]
59 feature_tree_operation: bool,
60
61 #[serde(default)]
64 keywords: bool,
65
66 #[serde(default)]
69 unlabeled_first: bool,
70
71 #[serde(default)]
73 args: HashMap<String, ArgMetadata>,
74}
75
76#[proc_macro_attribute]
77pub fn stdlib(attr: proc_macro::TokenStream, item: proc_macro::TokenStream) -> proc_macro::TokenStream {
78 do_output(do_stdlib(attr.into(), item.into()))
79}
80
81fn do_stdlib(
82 attr: proc_macro2::TokenStream,
83 item: proc_macro2::TokenStream,
84) -> Result<(proc_macro2::TokenStream, Vec<Error>), Error> {
85 let metadata = from_tokenstream(&attr)?;
86 do_stdlib_inner(metadata, attr, item)
87}
88
89fn do_output(res: Result<(proc_macro2::TokenStream, Vec<Error>), Error>) -> proc_macro::TokenStream {
90 match res {
91 Err(err) => err.to_compile_error().into(),
92 Ok((stdlib_docs, errors)) => {
93 let compiler_errors = errors.iter().map(|err| err.to_compile_error());
94
95 let output = quote! {
96 #stdlib_docs
97 #( #compiler_errors )*
98 };
99
100 output.into()
101 }
102 }
103}
104
105fn do_stdlib_inner(
106 metadata: StdlibMetadata,
107 _attr: proc_macro2::TokenStream,
108 item: proc_macro2::TokenStream,
109) -> Result<(proc_macro2::TokenStream, Vec<Error>), Error> {
110 let ast: ItemFnForSignature = syn::parse2(item.clone())?;
111
112 let mut errors = Vec::new();
113
114 if ast.sig.constness.is_some() {
115 errors.push(Error::new_spanned(
116 &ast.sig.constness,
117 "stdlib functions may not be const functions",
118 ));
119 }
120
121 if ast.sig.unsafety.is_some() {
122 errors.push(Error::new_spanned(
123 &ast.sig.unsafety,
124 "stdlib functions may not be unsafe",
125 ));
126 }
127
128 if ast.sig.abi.is_some() {
129 errors.push(Error::new_spanned(
130 &ast.sig.abi,
131 "stdlib functions may not use an alternate ABI",
132 ));
133 }
134
135 if !ast.sig.generics.params.is_empty() {
136 if ast.sig.generics.params.iter().any(|generic_type| match generic_type {
137 syn::GenericParam::Lifetime(_) => false,
138 syn::GenericParam::Type(_) => true,
139 syn::GenericParam::Const(_) => true,
140 }) {
141 errors.push(Error::new_spanned(
142 &ast.sig.generics,
143 "Stdlib functions may not be generic over types or constants, only lifetimes.",
144 ));
145 }
146 }
147
148 if ast.sig.variadic.is_some() {
149 errors.push(Error::new_spanned(&ast.sig.variadic, "no language C here"));
150 }
151
152 let name = metadata.name;
153
154 let name_cleaned = name.strip_suffix("2d").unwrap_or(name.as_str());
157 let name_cleaned = name.strip_suffix("3d").unwrap_or(name_cleaned);
158 if !name_cleaned.is_camel_case() {
159 errors.push(Error::new_spanned(
160 &ast.sig.ident,
161 format!("stdlib function names must be in camel case: `{}`", name),
162 ));
163 }
164
165 let name_ident = format_ident!("{}", name.to_case(convert_case::Case::UpperCamel));
166 let name_str = name.to_string();
167
168 let fn_name = &ast.sig.ident;
169 let fn_name_str = fn_name.to_string().replace("inner_", "");
170 let fn_name_ident = format_ident!("{}", fn_name_str);
171 let boxed_fn_name_ident = format_ident!("boxed_{}", fn_name_str);
172 let _visibility = &ast.vis;
173
174 let doc_info = extract_doc_from_attrs(&ast.attrs);
175 let comment_text = {
176 let mut buf = String::new();
177 buf.push_str("Std lib function: ");
178 buf.push_str(&name_str);
179 if let Some(s) = &doc_info.summary {
180 buf.push_str("\n");
181 buf.push_str(&s);
182 }
183 if let Some(s) = &doc_info.description {
184 buf.push_str("\n");
185 buf.push_str(&s);
186 }
187 buf
188 };
189 let description_doc_comment = quote! {
190 #[doc = #comment_text]
191 };
192
193 let summary = if let Some(summary) = doc_info.summary {
194 quote! { #summary }
195 } else {
196 quote! { "" }
197 };
198 let description = if let Some(description) = doc_info.description {
199 quote! { #description }
200 } else {
201 quote! { "" }
202 };
203
204 let cb = doc_info.code_blocks.clone();
205 let code_blocks = if !cb.is_empty() {
206 quote! {
207 let code_blocks = vec![#(#cb),*];
208 code_blocks.iter().map(|cb| {
209 let program = crate::Program::parse_no_errs(cb).unwrap();
210
211 let mut options: crate::parsing::ast::types::FormatOptions = Default::default();
212 options.insert_final_newline = false;
213 program.ast.recast(&options, 0)
214 }).collect::<Vec<String>>()
215 }
216 } else {
217 errors.push(Error::new_spanned(
218 &ast.sig,
219 "stdlib functions must have at least one code block",
220 ));
221
222 quote! { vec![] }
223 };
224
225 for code_block in doc_info.code_blocks.iter() {
227 if !code_block.contains(&name) {
228 errors.push(Error::new_spanned(
229 &ast.sig,
230 format!(
231 "stdlib functions must have the function name `{}` in the code block",
232 name
233 ),
234 ));
235 }
236 }
237
238 let test_code_blocks = doc_info
239 .code_blocks
240 .iter()
241 .enumerate()
242 .map(|(index, code_block)| generate_code_block_test(&fn_name_str, code_block, index))
243 .collect::<Vec<_>>();
244
245 let tags = metadata
246 .tags
247 .iter()
248 .map(|tag| {
249 quote! { #tag.to_string() }
250 })
251 .collect::<Vec<_>>();
252
253 let deprecated = if metadata.deprecated {
254 quote! { true }
255 } else {
256 quote! { false }
257 };
258
259 let unpublished = if metadata.unpublished {
260 quote! { true }
261 } else {
262 quote! { false }
263 };
264
265 let feature_tree_operation = if metadata.feature_tree_operation {
266 quote! { true }
267 } else {
268 quote! { false }
269 };
270
271 let uses_keyword_arguments = if metadata.keywords {
272 quote! { true }
273 } else {
274 quote! { false }
275 };
276
277 let docs_crate = get_crate(None);
278
279 let mut arg_types = Vec::new();
285 for (i, arg) in ast.sig.inputs.iter().enumerate() {
286 let arg_name = match arg {
288 syn::FnArg::Receiver(pat) => {
289 let span = pat.self_token.span.unwrap();
290 span.source_text().unwrap().to_string()
291 }
292 syn::FnArg::Typed(pat) => match &*pat.pat {
293 syn::Pat::Ident(ident) => ident.ident.to_string(),
294 _ => {
295 errors.push(Error::new_spanned(
296 &pat.pat,
297 "stdlib functions may not use destructuring patterns",
298 ));
299 continue;
300 }
301 },
302 }
303 .trim_start_matches('_')
304 .to_string();
305
306 let ty = match arg {
307 syn::FnArg::Receiver(pat) => pat.ty.as_ref().into_token_stream(),
308 syn::FnArg::Typed(pat) => pat.ty.as_ref().into_token_stream(),
309 };
310
311 let (ty_string, ty_ident) = clean_ty_string(ty.to_string().as_str());
312
313 let ty_string = rust_type_to_openapi_type(&ty_string);
314 let required = !ty_ident.to_string().starts_with("Option <");
315 let arg_meta = metadata.args.get(&arg_name);
316 let description = if let Some(s) = arg_meta.map(|arg| &arg.docs) {
317 quote! { #s }
318 } else if metadata.keywords && ty_string != "Args" && ty_string != "ExecState" {
319 errors.push(Error::new_spanned(
320 &arg,
321 "Argument was not documented in the args block",
322 ));
323 continue;
324 } else {
325 quote! { String::new() }
326 };
327 let include_in_snippet = required || arg_meta.map(|arg| arg.include_in_snippet).unwrap_or_default();
328 let label_required = !(i == 0 && metadata.unlabeled_first);
329 let camel_case_arg_name = to_camel_case(&arg_name);
330 if ty_string != "ExecState" && ty_string != "Args" {
331 let schema = quote! {
332 #docs_crate::cleanup_number_tuples_root(generator.root_schema_for::<#ty_ident>())
333 };
334 arg_types.push(quote! {
335 #docs_crate::StdLibFnArg {
336 name: #camel_case_arg_name.to_string(),
337 type_: #ty_string.to_string(),
338 schema: #schema,
339 required: #required,
340 label_required: #label_required,
341 description: #description.to_string(),
342 include_in_snippet: #include_in_snippet,
343 }
344 });
345 }
346 }
347
348 let return_type_inner = match &ast.sig.output {
349 syn::ReturnType::Default => quote! { () },
350 syn::ReturnType::Type(_, ty) => {
351 match &**ty {
353 syn::Type::Path(syn::TypePath { path, .. }) => {
354 let path = &path.segments;
355 if path.len() == 1 {
356 let seg = &path[0];
357 if seg.ident == "Result" {
358 if let syn::PathArguments::AngleBracketed(syn::AngleBracketedGenericArguments {
359 args,
360 ..
361 }) = &seg.arguments
362 {
363 if args.len() == 2 || args.len() == 1 {
364 let mut args = args.iter();
365 let ok = args.next().unwrap();
366 if let syn::GenericArgument::Type(ty) = ok {
367 let ty = unbox(ty.clone());
368 quote! { #ty }
369 } else {
370 quote! { () }
371 }
372 } else {
373 quote! { () }
374 }
375 } else {
376 quote! { () }
377 }
378 } else {
379 let ty = unbox(*ty.clone());
380 quote! { #ty }
381 }
382 } else {
383 quote! { () }
384 }
385 }
386 _ => {
387 quote! { () }
388 }
389 }
390 }
391 };
392
393 let ret_ty_string = return_type_inner.to_string().replace(' ', "");
394 let return_type = if !ret_ty_string.is_empty() || ret_ty_string != "()" {
395 let ret_ty_string = rust_type_to_openapi_type(&ret_ty_string);
396 quote! {
397 let schema = #docs_crate::cleanup_number_tuples_root(generator.root_schema_for::<#return_type_inner>());
398 Some(#docs_crate::StdLibFnArg {
399 name: "".to_string(),
400 type_: #ret_ty_string.to_string(),
401 schema,
402 required: true,
403 label_required: true,
404 description: String::new(),
405 include_in_snippet: true,
406 })
407 }
408 } else {
409 quote! {
410 None
411 }
412 };
413
414 let span = ast.sig.ident.span();
420 let const_struct = quote_spanned! {span=>
421 pub(crate) const #name_ident: #name_ident = #name_ident {};
422 };
423
424 let test_mod_name = format_ident!("test_examples_{}", fn_name_str);
425
426 let stream = quote! {
429 #[cfg(test)]
430 mod #test_mod_name {
431 #(#test_code_blocks)*
432 }
433
434 #[allow(non_camel_case_types, missing_docs)]
436 #description_doc_comment
437 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, schemars::JsonSchema, ts_rs::TS)]
438 #[ts(export)]
439 pub(crate) struct #name_ident {}
440 #[allow(non_upper_case_globals, missing_docs)]
442 #description_doc_comment
443 #const_struct
444
445 fn #boxed_fn_name_ident(
446 exec_state: &mut crate::execution::ExecState,
447 args: crate::std::Args,
448 ) -> std::pin::Pin<
449 Box<dyn std::future::Future<Output = anyhow::Result<crate::execution::KclValue, crate::errors::KclError>> + Send + '_>,
450 > {
451 Box::pin(#fn_name_ident(exec_state, args))
452 }
453
454 impl #docs_crate::StdLibFn for #name_ident
455 {
456 fn name(&self) -> String {
457 #name_str.to_string()
458 }
459
460 fn summary(&self) -> String {
461 #summary.to_string()
462 }
463
464 fn description(&self) -> String {
465 #description.to_string()
466 }
467
468 fn tags(&self) -> Vec<String> {
469 vec![#(#tags),*]
470 }
471
472 fn keyword_arguments(&self) -> bool {
473 #uses_keyword_arguments
474 }
475
476 fn args(&self, inline_subschemas: bool) -> Vec<#docs_crate::StdLibFnArg> {
477 let mut settings = schemars::gen::SchemaSettings::openapi3();
478 settings.inline_subschemas = inline_subschemas;
480 let mut generator = schemars::gen::SchemaGenerator::new(settings);
481
482 vec![#(#arg_types),*]
483 }
484
485 fn return_value(&self, inline_subschemas: bool) -> Option<#docs_crate::StdLibFnArg> {
486 let mut settings = schemars::gen::SchemaSettings::openapi3();
487 settings.inline_subschemas = inline_subschemas;
489 let mut generator = schemars::gen::SchemaGenerator::new(settings);
490
491 #return_type
492 }
493
494 fn unpublished(&self) -> bool {
495 #unpublished
496 }
497
498 fn deprecated(&self) -> bool {
499 #deprecated
500 }
501
502 fn feature_tree_operation(&self) -> bool {
503 #feature_tree_operation
504 }
505
506 fn examples(&self) -> Vec<String> {
507 #code_blocks
508 }
509
510 fn std_lib_fn(&self) -> crate::std::StdFn {
511 #boxed_fn_name_ident
512 }
513
514 fn clone_box(&self) -> Box<dyn #docs_crate::StdLibFn> {
515 Box::new(self.clone())
516 }
517 }
518
519 #item
520 };
521
522 if !errors.is_empty() {
524 errors.insert(0, Error::new_spanned(&ast.sig, ""));
525 }
526
527 Ok((stream, errors))
528}
529
530#[allow(dead_code)]
531fn to_compile_errors(errors: Vec<syn::Error>) -> proc_macro2::TokenStream {
532 let compile_errors = errors.iter().map(syn::Error::to_compile_error);
533 quote!(#(#compile_errors)*)
534}
535
536fn get_crate(var: Option<String>) -> proc_macro2::TokenStream {
537 if let Some(s) = var {
538 if let Ok(ts) = syn::parse_str(s.as_str()) {
539 return ts;
540 }
541 }
542 quote!(crate::docs)
543}
544
545#[derive(Debug)]
546struct DocInfo {
547 pub summary: Option<String>,
548 pub description: Option<String>,
549 pub code_blocks: Vec<String>,
550}
551
552fn extract_doc_from_attrs(attrs: &[syn::Attribute]) -> DocInfo {
553 let doc = syn::Ident::new("doc", proc_macro2::Span::call_site());
554 let raw_lines = attrs.iter().flat_map(|attr| {
555 if let syn::Meta::NameValue(nv) = &attr.meta {
556 if nv.path.is_ident(&doc) {
557 if let syn::Expr::Lit(syn::ExprLit {
558 lit: syn::Lit::Str(s), ..
559 }) = &nv.value
560 {
561 return normalize_comment_string(s.value());
562 }
563 }
564 }
565 Vec::new()
566 });
567
568 let mut code_blocks: Vec<String> = Vec::new();
570 let mut code_block: Option<String> = None;
571 let mut parsed_lines = Vec::new();
572 for line in raw_lines {
573 if line.starts_with("```") {
574 if let Some(ref inner_code_block) = code_block {
575 code_blocks.push(inner_code_block.trim().to_string());
576 code_block = None;
577 } else {
578 code_block = Some(String::new());
579 }
580
581 continue;
582 }
583 if let Some(ref mut code_block) = code_block {
584 code_block.push_str(&line);
585 code_block.push('\n');
586 } else {
587 parsed_lines.push(line);
588 }
589 }
590
591 if let Some(code_block) = code_block {
592 code_blocks.push(code_block.trim().to_string());
593 }
594
595 let mut summary = None;
596 let mut description: Option<String> = None;
597 for line in parsed_lines {
598 if line.is_empty() {
599 if let Some(desc) = &mut description {
600 if !desc.is_empty() && !desc.ends_with('\n') {
602 if desc.ends_with(' ') {
603 desc.pop().unwrap();
604 }
605 desc.push_str("\n\n");
606 }
607 } else if summary.is_some() {
608 description = Some(String::new());
609 }
610 continue;
611 }
612
613 if let Some(desc) = &mut description {
614 desc.push_str(&line);
615 desc.push(' ');
617 continue;
618 }
619
620 if summary.is_none() {
621 summary = Some(String::new());
622 }
623 match &mut summary {
624 Some(summary) => {
625 summary.push_str(&line);
626 summary.push(' ');
628 }
629 None => unreachable!(),
630 }
631 }
632
633 if let Some(s) = &mut summary {
635 while s.ends_with(' ') || s.ends_with('\n') {
636 s.pop().unwrap();
637 }
638
639 if s.is_empty() {
640 summary = None;
641 }
642 }
643
644 if let Some(d) = &mut description {
645 while d.ends_with(' ') || d.ends_with('\n') {
646 d.pop().unwrap();
647 }
648
649 if d.is_empty() {
650 description = None;
651 }
652 }
653
654 DocInfo {
655 summary,
656 description,
657 code_blocks,
658 }
659}
660
661fn normalize_comment_string(s: String) -> Vec<String> {
662 s.split('\n')
663 .map(|s| {
664 s.strip_prefix(' ').unwrap_or(s).trim_end().to_owned()
668 })
669 .collect()
670}
671
672struct ItemFnForSignature {
675 pub attrs: Vec<Attribute>,
676 pub vis: Visibility,
677 pub sig: Signature,
678 pub _block: proc_macro2::TokenStream,
679}
680
681impl Parse for ItemFnForSignature {
682 fn parse(input: ParseStream) -> syn::parse::Result<Self> {
683 let attrs = input.call(Attribute::parse_outer)?;
684 let vis: Visibility = input.parse()?;
685 let sig: Signature = input.parse()?;
686 let block = input.parse()?;
687 Ok(ItemFnForSignature {
688 attrs,
689 vis,
690 sig,
691 _block: block,
692 })
693 }
694}
695
696fn clean_ty_string(t: &str) -> (String, proc_macro2::TokenStream) {
697 let mut ty_string = t
698 .replace("& 'a", "")
699 .replace('&', "")
700 .replace("mut", "")
701 .replace("< 'a >", "")
702 .replace(' ', "");
703 if ty_string.starts_with("ExecState") {
704 ty_string = "ExecState".to_string();
705 }
706 if ty_string.starts_with("Args") {
707 ty_string = "Args".to_string();
708 }
709 let ty_string = ty_string.trim().to_string();
710 let ty_ident = if ty_string.starts_with("Vec<") {
711 let ty_string = ty_string.trim_start_matches("Vec<").trim_end_matches('>');
712 let (_, ty_ident) = clean_ty_string(&ty_string);
713 quote! {
714 Vec<#ty_ident>
715 }
716 } else if ty_string.starts_with("kittycad::types::") {
717 let ty_string = ty_string.trim_start_matches("kittycad::types::").trim_end_matches('>');
718 let ty_ident = format_ident!("{}", ty_string);
719 quote! {
720 kittycad::types::#ty_ident
721 }
722 } else if ty_string.starts_with("Option<") {
723 let ty_string = ty_string.trim_start_matches("Option<").trim_end_matches('>');
724 let (_, ty_ident) = clean_ty_string(&ty_string);
725 quote! {
726 Option<#ty_ident>
727 }
728 } else if let Some((inner_array_type, num)) = parse_array_type(&ty_string) {
729 let ty_string = inner_array_type.to_owned();
730 let (_, ty_ident) = clean_ty_string(&ty_string);
731 quote! {
732 [#ty_ident; #num]
733 }
734 } else if ty_string.starts_with("Box<") {
735 let ty_string = ty_string.trim_start_matches("Box<").trim_end_matches('>');
736 let (_, ty_ident) = clean_ty_string(&ty_string);
737 quote! {
738 #ty_ident
739 }
740 } else {
741 let ty_ident = format_ident!("{}", ty_string);
742 quote! {
743 #ty_ident
744 }
745 };
746
747 (ty_string, ty_ident)
748}
749
750fn rust_type_to_openapi_type(t: &str) -> String {
751 let mut t = t.to_string();
752 if t.starts_with("Vec<") {
755 t = t.replace("Vec<", "[").replace('>', "]");
756 }
757 if t.starts_with("Box<") {
758 t = t.replace("Box<", "").replace('>', "");
759 }
760 if t.starts_with("Option<") {
761 t = t.replace("Option<", "").replace('>', "");
762 }
763 if let Some((inner_type, _length)) = parse_array_type(&t) {
764 t = format!("[{inner_type}]")
765 }
766
767 if t == "f64" {
768 return "number".to_string();
769 } else if t == "u32" {
770 return "integer".to_string();
771 } else if t == "str" {
772 return "string".to_string();
773 } else {
774 return t.replace("f64", "number").to_string();
775 }
776}
777
778fn parse_array_type(type_name: &str) -> Option<(&str, usize)> {
779 static RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\[([a-zA-Z0-9<>]+); ?(\d+)\]").unwrap());
780 let cap = RE.captures(type_name)?;
781 let inner_type = cap.get(1)?;
782 let length = cap.get(2)?.as_str().parse().ok()?;
783 Some((inner_type.as_str(), length))
784}
785
786fn generate_code_block_test(fn_name: &str, code_block: &str, index: usize) -> proc_macro2::TokenStream {
789 let test_name = format_ident!("kcl_test_example_{}{}", fn_name, index);
790 let test_name_mock = format_ident!("test_mock_example_{}{}", fn_name, index);
791 let output_test_name_str = format!("serial_test_example_{}{}", fn_name, index);
792
793 quote! {
794 #[tokio::test(flavor = "multi_thread")]
795 async fn #test_name_mock() -> miette::Result<()> {
796 let program = crate::Program::parse_no_errs(#code_block).unwrap();
797 let ctx = crate::ExecutorContext {
798 engine: std::sync::Arc::new(Box::new(crate::engine::conn_mock::EngineConnection::new().await.unwrap())),
799 fs: std::sync::Arc::new(crate::fs::FileManager::new()),
800 stdlib: std::sync::Arc::new(crate::std::StdLib::new()),
801 settings: Default::default(),
802 context_type: crate::execution::ContextType::Mock,
803 };
804
805 if let Err(e) = ctx.run(&program, &mut crate::execution::ExecState::new(&ctx)).await {
806 return Err(miette::Report::new(crate::errors::Report {
807 error: e.error,
808 filename: format!("{}{}", #fn_name, #index),
809 kcl_source: #code_block.to_string(),
810 }));
811 }
812 Ok(())
813 }
814
815 #[tokio::test(flavor = "multi_thread", worker_threads = 5)]
816 async fn #test_name() -> miette::Result<()> {
817 let code = #code_block;
818 let result = match crate::test_server::execute_and_snapshot(code, crate::settings::types::UnitLength::Mm, None).await {
820 Err(crate::errors::ExecError::Kcl(e)) => {
821 return Err(miette::Report::new(crate::errors::Report {
822 error: e.error,
823 filename: format!("{}{}", #fn_name, #index),
824 kcl_source: #code_block.to_string(),
825 }));
826 }
827 Err(other_err)=> panic!("{}", other_err),
828 Ok(img) => img,
829 };
830 twenty_twenty::assert_image(&format!("tests/outputs/{}.png", #output_test_name_str), &result, 0.99);
831 Ok(())
832 }
833 }
834}