feat(napi-derive): add ts_arg_type attribute to override individual args on functions (#1192)

crates/backend/src/ast.rs

@@ -6,7 +6,7 @@ pub struct NapiFn {
   pub name: Ident,
   pub js_name: String,
   pub attrs: Vec<Attribute>,
-  pub args: Vec<NapiFnArgKind>,
+  pub args: Vec<NapiFnArg>,
   pub ret: Option<syn::Type>,
   pub is_ret_result: bool,
   pub is_async: bool,
@@ -31,6 +31,23 @@ pub struct CallbackArg {
   pub ret: Option<syn::Type>,
 }
 
+#[derive(Debug, Clone)]
+pub struct NapiFnArg {
+  pub kind: NapiFnArgKind,
+  pub ts_arg_type: Option<String>,
+}
+
+impl NapiFnArg {
+  /// if type was overridden with `#[napi(ts_arg_type = "...")]` use that instead
+  pub fn use_overridden_type_or(&self, default: impl FnOnce() -> String) -> String {
+    self
+      .ts_arg_type
+      .as_ref()
+      .map(|ts| ts.clone())
+      .unwrap_or_else(default)
+  }
+}
+
 #[derive(Debug, Clone)]
 pub enum NapiFnArgKind {
   PatType(Box<syn::PatType>),

crates/backend/src/codegen/fn.rs

@@ -121,7 +121,7 @@ impl NapiFn {
       let i = i - skipped_arg_count;
       let ident = Ident::new(&format!("arg{}", i), Span::call_site());
 
-      match arg {
+      match &arg.kind {
         NapiFnArgKind::PatType(path) => {
           if &path.ty.to_token_stream().to_string() == "Env" {
             args.push(quote! { napi::bindgen_prelude::Env::from(env) });

crates/backend/src/typegen/fn.rs

@@ -119,7 +119,7 @@ impl NapiFn {
       self
         .args
         .iter()
-        .filter_map(|arg| match arg {
+        .filter_map(|arg| match &arg.kind {
           crate::NapiFnArgKind::PatType(path) => {
             let ty_string = path.ty.to_token_stream().to_string();
             if ty_string == "Env" {
@@ -142,8 +142,10 @@ impl NapiFn {
             if let Pat::Ident(i) = path.pat.as_mut() {
               i.mutability = None;
             }
-            let arg = path.pat.to_token_stream().to_string().to_case(Case::Camel);
+
             let (ts_type, is_optional) = ty_to_ts_type(&path.ty, false, false);
+            let ts_type = arg.use_overridden_type_or(|| ts_type);
+            let arg = path.pat.to_token_stream().to_string().to_case(Case::Camel);
 
             Some(FnArg {
               arg,
@@ -152,8 +154,8 @@ impl NapiFn {
             })
           }
           crate::NapiFnArgKind::Callback(cb) => {
+            let ts_type = arg.use_overridden_type_or(|| gen_callback_type(cb));
             let arg = cb.pat.to_token_stream().to_string().to_case(Case::Camel);
-            let ts_type = gen_callback_type(cb);
 
             Some(FnArg {
               arg,

crates/macro/src/parser/mod.rs

@@ -10,14 +10,14 @@ use attrs::{BindgenAttr, BindgenAttrs};
 use convert_case::{Case, Casing};
 use napi_derive_backend::{
   BindgenResult, CallbackArg, Diagnostic, FnKind, FnSelf, Napi, NapiConst, NapiEnum,
-  NapiEnumVariant, NapiFn, NapiFnArgKind, NapiImpl, NapiItem, NapiStruct, NapiStructField,
+  NapiEnumVariant, NapiFn, NapiFnArg, NapiFnArgKind, NapiImpl, NapiItem, NapiStruct,
-  NapiStructKind,
+  NapiStructField, NapiStructKind,
 };
-use proc_macro2::{Ident, TokenStream, TokenTree};
+use proc_macro2::{Ident, Span, TokenStream, TokenTree};
 use quote::ToTokens;
 use syn::ext::IdentExt;
 use syn::parse::{Parse, ParseStream, Result as SynResult};
-use syn::{Attribute, PathSegment, Signature, Type, Visibility};
+use syn::{Attribute, Meta, NestedMeta, PatType, PathSegment, Signature, Type, Visibility};
 
 use crate::parser::attrs::{check_recorded_struct_for_impl, record_struct};
 
@@ -155,6 +155,60 @@ pub trait ParseNapi {
   fn parse_napi(&mut self, tokens: &mut TokenStream, opts: BindgenAttrs) -> BindgenResult<Napi>;
 }
 
+/// This function does a few things:
+/// - parses the tokens for the given argument `p` to find the `#[napi(ts_arg_type = "MyType")]`
+///   attribute and return the manually overridden type.
+/// - If both the `ts_args_type` override and the `ts_arg_type` override are present, bail
+///   since it should only allow one at a time.
+/// - Bails if it finds the `#[napi...]` attribute but it has the wrong data.
+/// - Removes the attribute from the output token stream so this
+///   `pub fn add(u: u32, #[napi(ts_arg_type = "MyType")] f: String)`
+///    turns into
+///   `pub fn add(u: u32, f: String)`
+///    otherwise it won't compile
+fn find_ts_arg_type_and_remove_attribute(
+  p: &mut PatType,
+  ts_args_type: Option<&(&str, Span)>,
+) -> BindgenResult<Option<String>> {
+  for (idx, attr) in p.attrs.iter().enumerate() {
+    if let Ok(Meta::List(meta_list)) = attr.parse_meta() {
+      if meta_list.path.get_ident() != Some(&format_ident!("napi")) {
+        // If this attribute is not for `napi` ignore it.
+        continue;
+      }
+
+      if let Some((ts_args_type, _)) = ts_args_type {
+        bail_span!(
+          meta_list,
+          "Found a 'ts_args_type'=\"{}\" override. Cannot use 'ts_arg_type' at the same time since they are mutually exclusive.",
+          ts_args_type
+        );
+      }
+
+      let nested = meta_list.nested.first();
+
+      let nm = if let Some(NestedMeta::Meta(Meta::NameValue(nm))) = nested {
+        nm
+      } else {
+        bail_span!(meta_list.nested, "Expected Name Value");
+      };
+
+      if Some(&format_ident!("ts_arg_type")) != nm.path.get_ident() {
+        bail_span!(nm.path, "Did not find 'ts_arg_type'");
+      }
+
+      if let syn::Lit::Str(lit) = &nm.lit {
+        p.attrs.remove(idx);
+        return Ok(Some(lit.value()));
+      } else {
+        bail_span!(nm.lit, "Expected a string literal");
+      }
+    }
+  }
+
+  Ok(None)
+}
+
 fn get_ty(mut ty: &syn::Type) -> &syn::Type {
   while let syn::Type::Group(g) = ty {
     ty = &g.elem;
@@ -462,7 +516,7 @@ fn extract_fn_closure_generics(
 }
 
 fn napi_fn_from_decl(
-  sig: Signature,
+  sig: &mut Signature,
   opts: &BindgenAttrs,
   attrs: Vec<Attribute>,
   vis: Visibility,
@@ -473,36 +527,48 @@ fn napi_fn_from_decl(
   let syn::Signature {
     ident,
     asyncness,
-    inputs,
     output,
     generics,
     ..
-  } = sig;
+  } = sig.clone();
 
   let mut fn_self = None;
   let callback_traits = extract_fn_closure_generics(&generics)?;
 
-  let args = inputs
+  let args = sig
-    .into_iter()
+    .inputs
+    .iter_mut()
     .filter_map(|arg| match arg {
-      syn::FnArg::Typed(mut p) => {
+      syn::FnArg::Typed(ref mut p) => {
+        let ts_arg_type = find_ts_arg_type_and_remove_attribute(p, opts.ts_args_type().as_ref())
+          .unwrap_or_else(|e| {
+            errors.push(e);
+            None
+          });
+
         let ty_str = p.ty.to_token_stream().to_string();
         if let Some(path_arguments) = callback_traits.get(&ty_str) {
           match extract_callback_trait_types(path_arguments) {
-            Ok((fn_args, fn_ret)) => Some(NapiFnArgKind::Callback(Box::new(CallbackArg {
+            Ok((fn_args, fn_ret)) => Some(NapiFnArg {
-              pat: p.pat,
+              kind: NapiFnArgKind::Callback(Box::new(CallbackArg {
-              args: fn_args,
+                pat: p.pat.clone(),
-              ret: fn_ret,
+                args: fn_args,
-            }))),
+                ret: fn_ret,
+              })),
+              ts_arg_type,
+            }),
             Err(e) => {
               errors.push(e);
               None
             }
           }
         } else {
-          let ty = replace_self(*p.ty, parent);
+          let ty = replace_self(p.ty.as_ref().clone(), parent);
           p.ty = Box::new(ty);
-          Some(NapiFnArgKind::PatType(Box::new(p)))
+          Some(NapiFnArg {
+            kind: NapiFnArgKind::PatType(Box::new(p.clone())),
+            ts_arg_type,
+          })
         }
       }
       syn::FnArg::Receiver(r) => {
@@ -638,8 +704,10 @@ impl ParseNapi for syn::ItemFn {
         "#[napi] can't be applied to a function with #[napi(ts_type)]"
       );
     }
+    let napi = self.convert_to_ast(opts);
     self.to_tokens(tokens);
-    self.convert_to_ast(opts)
+
+    napi
   }
 }
 impl ParseNapi for syn::ItemStruct {
@@ -734,7 +802,7 @@ fn fn_kind(opts: &BindgenAttrs) -> FnKind {
 impl ConvertToAST for syn::ItemFn {
   fn convert_to_ast(&mut self, opts: BindgenAttrs) -> BindgenResult<Napi> {
     let func = napi_fn_from_decl(
-      self.sig.clone(),
+      &mut self.sig,
       &opts,
       self.attrs.clone(),
       self.vis.clone(),
@@ -912,7 +980,7 @@ impl ConvertToAST for syn::ItemImpl {
         }
 
         let func = napi_fn_from_decl(
-          method.sig.clone(),
+          &mut method.sig,
           &opts,
           method.attrs.clone(),
           vis,

examples/napi/__test__/typegen.spec.ts.md

@@ -105,6 +105,8 @@ Generated by [AVA](https://avajs.dev).
     export function validateString(s: string): string␊
     export function validateSymbol(s: symbol): boolean␊
     export function tsRename(a: { foo: number }): string[]␊
+    export function overrideIndividualArgOnFunction(notOverridden: string, f: () => string, notOverridden2: number): string␊
+    export function overrideIndividualArgOnFunctionWithCbArg(callback: (town: string, name?: string | undefined | null) => string, notOverridden: number): object␊
     export function xxh64Alias(input: Buffer): bigint␊
     export function getMapping(): Record<string, number>␊
     export function sumMapping(nums: Record<string, number>): number␊
@@ -208,6 +210,7 @@ Generated by [AVA](https://avajs.dev).
        */␊
       returnOtherClass(): Dog␊
       returnOtherClassWithCustomConstructor(): Bird␊
+      overrideIndividualArgOnMethod(normalTy: string, overriddenTy: {n: string}): Bird␊
     }␊
     export class Dog {␊
       name: string␊

examples/napi/__test__/values.spec.ts

@@ -95,6 +95,8 @@ import {
   callbackReturnPromise,
   returnEitherClass,
   eitherFromOption,
+  overrideIndividualArgOnFunction,
+  overrideIndividualArgOnFunctionWithCbArg,
 } from '../'
 
 test('export const', (t) => {
@@ -157,6 +159,10 @@ test('class', (t) => {
   t.is(dog.name, '可乐')
   t.deepEqual(dog.returnOtherClass(), new Dog('Doge'))
   t.deepEqual(dog.returnOtherClassWithCustomConstructor(), new Bird('parrot'))
+  t.is(
+    dog.overrideIndividualArgOnMethod('Jafar', { n: 'Iago' }).name,
+    'Jafar-Iago',
+  )
   t.is(dog.returnOtherClassWithCustomConstructor().getCount(), 1234)
   t.is(dog.type, Kind.Dog)
   dog.type = Kind.Cat
@@ -321,6 +327,20 @@ test('function ts type override', (t) => {
   t.deepEqual(tsRename({ foo: 1, bar: 2, baz: 2 }), ['foo', 'bar', 'baz'])
 })
 
+test('function individual ts arg type override', (t) => {
+  t.is(
+    overrideIndividualArgOnFunction('someStr', () => 'anotherStr', 42),
+    'oia: someStr-42-anotherStr',
+  )
+  t.deepEqual(
+    overrideIndividualArgOnFunctionWithCbArg(
+      (town, opt) => `im: ${town}-${opt}`,
+      89,
+    ),
+    'im: World(89)-null',
+  )
+})
+
 test('option object', (t) => {
   t.notThrows(() => receiveAllOptionalObject())
   t.notThrows(() => receiveAllOptionalObject({}))

examples/napi/index.d.ts

@@ -95,6 +95,8 @@ export function validatePromise(p: Promise<number>): Promise<number>
 export function validateString(s: string): string
 export function validateSymbol(s: symbol): boolean
 export function tsRename(a: { foo: number }): string[]
+export function overrideIndividualArgOnFunction(notOverridden: string, f: () => string, notOverridden2: number): string
+export function overrideIndividualArgOnFunctionWithCbArg(callback: (town: string, name?: string | undefined | null) => string, notOverridden: number): object
 export function xxh64Alias(input: Buffer): bigint
 export function getMapping(): Record<string, number>
 export function sumMapping(nums: Record<string, number>): number
@@ -198,6 +200,7 @@ export class Animal {
    */
   returnOtherClass(): Dog
   returnOtherClassWithCustomConstructor(): Bird
+  overrideIndividualArgOnMethod(normalTy: string, overriddenTy: {n: string}): Bird
 }
 export class Dog {
   name: string

examples/napi/src/class.rs

@@ -85,6 +85,18 @@ impl Animal {
   pub fn return_other_class_with_custom_constructor(&self) -> Bird {
     Bird::new("parrot".to_owned())
   }
+
+  #[napi]
+  pub fn override_individual_arg_on_method(
+    &self,
+    normal_ty: String,
+    #[napi(ts_arg_type = "{n: string}")] overridden_ty: napi::JsObject,
+  ) -> Bird {
+    let obj = overridden_ty.coerce_to_object().unwrap();
+    let the_n: Option<String> = obj.get("n").unwrap();
+
+    Bird::new(format!("{}-{}", normal_ty, the_n.unwrap()))
+  }
 }
 
 #[napi(constructor)]

examples/napi/src/fn_ts_override.rs

@@ -1,6 +1,36 @@
 use napi::bindgen_prelude::{Object, Result};
+use napi::JsFunction;
 
 #[napi(ts_args_type = "a: { foo: number }", ts_return_type = "string[]")]
 fn ts_rename(a: Object) -> Result<Object> {
   a.get_property_names()
 }
+
+#[napi]
+fn override_individual_arg_on_function(
+  not_overridden: String,
+  #[napi(ts_arg_type = "() => string")] f: JsFunction,
+  not_overridden2: u32,
+) -> String {
+  let u = f.call_without_args(None).unwrap();
+  let s = u
+    .coerce_to_string()
+    .unwrap()
+    .into_utf8()
+    .unwrap()
+    .as_str()
+    .unwrap()
+    .to_string();
+
+  format!("oia: {}-{}-{}", not_overridden, not_overridden2, s)
+}
+
+#[napi]
+fn override_individual_arg_on_function_with_cb_arg<
+  T: Fn(String, Option<String>) -> Result<Object>,
+>(
+  #[napi(ts_arg_type = "(town: string, name?: string | undefined | null) => string")] callback: T,
+  not_overridden: u32,
+) -> Result<Object> {
+  callback(format!("World({})", not_overridden), None)
+}