feat: improve function typing (#2)

Co-authored-by: Pooya Parsa <[email protected]>

playground/consts.ts

@@ -1,5 +1,7 @@
 export const defaultReference = `
-export function add (id: String, date = new Date(), append?: Boolean) {}
+export function sendMessage (message: string, date = new Date(), flash?: boolean): string {
+  return 'OK'
+}
 
 export const config = {
     name: 'default',

src/generator/dts.ts

@@ -1,5 +1,5 @@
-import type { Schema, JSType } from '../types'
-import { escapeKey, unique } from '../utils'
+import type { Schema, JSType, TypeDescriptor } from '../types'
+import { escapeKey, normalizeTypes } from '../utils'
 
 const TYPE_MAP: Record<JSType, string> = {
   array: 'any[]',
@@ -24,7 +24,8 @@ const SCHEMA_KEYS = [
   'type',
   'tags',
   'args',
-  'id'
+  'id',
+  'returns'
 ]
 
 export function generateTypes (schema: Schema, name: string = 'Untyped') {
@@ -42,12 +43,11 @@ function _genTypes (schema: Schema, spaces: string): string[] {
     } else {
       let type: string
       if (val.type === 'array') {
-        const _type = getTsType(val.items.type)
-        type = _type.includes('|') ? `(${_type})[]` : `${_type}[]`
+        type = `Array<${getTsType(val.items)}>`
       } else if (val.type === 'function') {
-        type = `(${genFunctionArgs(val.args)}) => any`
+        type = genFunctionType(val)
       } else {
-        type = getTsType(val.type)
+        type = getTsType(val)
       }
       buff.push(`${escapeKey(key)}: ${type},\n`)
     }
@@ -63,24 +63,34 @@ function _genTypes (schema: Schema, spaces: string): string[] {
   return buff.map(i => spaces + i)
 }
 
-function getTsType (type: JSType | JSType[]): string {
+function getTsType (type: TypeDescriptor | TypeDescriptor[]): string {
   if (Array.isArray(type)) {
-    return unique(type.map(t => getTsType(t))).join(' | ') || 'any'
+    return normalizeTypes(type.map(t => getTsType(t))).join('|') || 'any'
   }
-  return (type && TYPE_MAP[type]) || 'any'
+  if (!type || !type.type) {
+    return 'any'
+  }
+  if (Array.isArray(type.type)) {
+    return type.type.map(t => TYPE_MAP[t]).join('|')
+  }
+  if (type.type === 'array') {
+    return `Array<${getTsType(type.items)}>`
+  }
+  return TYPE_MAP[type.type] || type.type
+}
+
+export function genFunctionType (schema) {
+  return `(${genFunctionArgs(schema.args)}) => ${getTsType(schema.returns)}`
 }
 
 export function genFunctionArgs (args: Schema['args']) {
   return args.map((arg) => {
     let argStr = arg.name
-    if (arg.optional) {
+    if (arg.optional || arg.default) {
       argStr += '?'
     }
     if (arg.type) {
-      argStr += ': ' + arg.type
-    }
-    if (arg.default) {
-      argStr += ' = ' + arg.default
+      argStr += `: ${getTsType(arg)}`
     }
     return argStr
   }).join(', ')

src/generator/md.ts

@@ -1,5 +1,5 @@
 import type { Schema } from '../types'
-import { genFunctionArgs } from './dts'
+import { genFunctionType } from './dts'
 
 export function generateMarkdown (schema: Schema) {
   return _generateMarkdown(schema, '', '').join('\n')
@@ -32,7 +32,7 @@ export function _generateMarkdown (schema: Schema, title: string, level) {
 
   // Signuture (function)
   if (schema.type === 'function') {
-    lines.push('```ts', '(genFunctionArgs(schema.args)) => {}', '```', '')
+    lines.push('```ts', genFunctionType(schema), '```', '')
   }
 
   // Description

src/loader/babel.ts

@@ -1,10 +1,24 @@
 import type { PluginObj } from '@babel/core'
 import * as t from '@babel/types'
-import { Schema } from '../types'
+import { Schema, JSType, TypeDescriptor, FunctionArg } from '../types'
+import { normalizeTypes, mergedTypes, cachedFn } from '../utils'
+
+type GetCodeFn = (loc: t.SourceLocation) => string
 
 export default function babelPluginUntyped () {
   return <PluginObj>{
     visitor: {
+      VariableDeclaration (p) {
+        const declaration = p.node.declarations[0]
+        if (
+          t.isIdentifier(declaration.id) &&
+          (t.isFunctionExpression(declaration.init) || t.isArrowFunctionExpression(declaration.init))
+        ) {
+          const newDeclaration = t.functionDeclaration(declaration.id, declaration.init.params, declaration.init.body as t.BlockStatement)
+          newDeclaration.returnType = declaration.init.returnType
+          p.replaceWith(newDeclaration)
+        }
+      },
       ObjectProperty (p) {
         if (p.node.leadingComments) {
           const schema = parseJSDocs(
@@ -20,20 +34,20 @@ export default function babelPluginUntyped () {
             if (schemaProp && 'value' in schemaProp) {
               if (schemaProp.value.type === 'ObjectExpression') {
                 // Object has $schema
-                schemaProp.value.properties.push(...schemaToPropsAst(schema))
+                schemaProp.value.properties.push(...astify(schema).properties)
               } else {
                 // Object has $schema which is not an object
                 // SKIP
               }
             } else {
               // Object has not $schema
-              p.node.value.properties.unshift(buildObjectPropery('$schema', schemaToAst(schema)))
+              p.node.value.properties.unshift(astify({ $schema: schema }))
             }
           } else {
             // Literal value
             p.node.value = t.objectExpression([
               t.objectProperty(t.identifier('$default'), p.node.value),
-              t.objectProperty(t.identifier('$schema'), schemaToAst(schema))
+              t.objectProperty(t.identifier('$schema'), astify(schema))
             ])
           }
           p.node.leadingComments = []
@@ -46,42 +60,52 @@ export default function babelPluginUntyped () {
             .filter(c => c.type === 'CommentBlock')
             .map(c => c.value)
         )
-
         schema.type = 'function'
         schema.args = []
 
-        const code = this.file.code.split('\n')
-        const getCode = loc => code[loc.start.line - 1].slice(loc.start.column, loc.end.column).trim() || ''
+        const _getLines = cachedFn(() => this.file.code.split('\n'))
+        const getCode: GetCodeFn = loc => _getLines()[loc.start.line - 1].slice(loc.start.column, loc.end.column).trim() || ''
 
         // Extract arguments
         p.node.params.forEach((param, index) => {
           if (param.loc.end.line !== param.loc.start.line) {
             return null
           }
-          if (param.type !== 'AssignmentPattern' && param.type !== 'Identifier') {
+          if (!t.isAssignmentPattern(param) && !t.isIdentifier(param)) {
             return null
           }
-          const _param = param.type === 'AssignmentPattern' ? param.left : param
-          const arg = {
-            // @ts-ignore TODO
-            name: _param.name || ('arg' + index),
-            type: getCode(_param.loc).split(':').slice(1).join(':').trim() || undefined,
-            default: undefined,
-            // @ts-ignore TODO
-            optional: _param.optional || undefined
+          const lparam = (t.isAssignmentPattern(param) ? param.left : param) as t.Identifier
+          if (!t.isIdentifier(lparam)) {
+            return null
+          }
+          const arg: FunctionArg = {
+            name: lparam.name || ('arg' + index),
+            optional: lparam.optional || undefined
+          }
+
+          // Infer from type annotations
+          if (lparam.typeAnnotation) {
+            Object.assign(arg, mergedTypes(arg, inferAnnotationType(lparam.typeAnnotation, getCode)))
           }
 
+          // Infer type from default value
           if (param.type === 'AssignmentPattern') {
-            arg.default = getCode(param.right.loc)
+            Object.assign(arg, mergedTypes(arg, inferArgType(param.right, getCode)))
           }
           schema.args.push(arg)
         })
 
+        // Return type annotation
+        if (p.node.returnType?.type === 'TSTypeAnnotation') {
+          schema.returns = inferAnnotationType(p.node.returnType, getCode)
+        }
+
         // Replace function with it's meta
-        const schemaAst = t.objectExpression([
-          buildObjectPropery('$schema', t.objectExpression(schemaToPropsAst(schema)))
-        ])
-        p.replaceWith(t.variableDeclaration('const', [t.variableDeclarator(t.identifier(p.node.id.name), schemaAst)]))
+        p.replaceWith(t.variableDeclaration('const', [
+          t.variableDeclarator(
+            t.identifier(p.node.id.name), astify({ $schema: schema })
+          )
+        ]))
       }
     }
   }
@@ -116,7 +140,7 @@ function parseJSDocs (input: string | string[]): Schema {
   return schema
 }
 
-function valToAstLiteral (val: any) {
+function astify (val: any) {
   if (typeof val === 'string') {
     return t.stringLiteral(val)
   }
@@ -126,42 +150,92 @@ function valToAstLiteral (val: any) {
   if (typeof val === 'number') {
     return t.numericLiteral(val)
   }
-  return null
-}
-
-function buildObjectPropsAst (obj: any) {
-  const props = []
-  for (const key in obj) {
-    const astLiteral = valToAstLiteral(obj[key])
-    if (astLiteral) {
-      props.push(t.objectProperty(t.identifier(key), astLiteral))
-    }
+  if (val === null) {
+    return t.nullLiteral()
+  }
+  if (val === undefined) {
+    return t.identifier('undefined')
   }
-  return props
+  if (Array.isArray(val)) {
+    return t.arrayExpression(val.map(item => astify(item)))
+  }
+  return t.objectExpression(Object.getOwnPropertyNames(val)
+    .filter(key => val[key] !== undefined && val[key] !== null)
+    .map(key => t.objectProperty(t.identifier(key), astify(val[key])))
+  )
 }
 
-function buildObjectPropery (name, val) {
-  return t.objectProperty(t.identifier(name), val)
+const AST_JSTYPE_MAP: Partial<Record<t.Expression['type'], JSType>> = {
+  StringLiteral: 'string',
+  BooleanLiteral: 'boolean',
+  BigIntLiteral: 'bigint',
+  DecimalLiteral: 'number',
+  NumericLiteral: 'number',
+  ObjectExpression: 'object',
+  FunctionExpression: 'function',
+  ArrowFunctionExpression: 'function',
+  RegExpLiteral: 'RegExp' as JSType
 }
 
-function schemaToPropsAst (schema: Schema) {
-  const props = buildObjectPropsAst(schema)
-
-  if (schema.args) {
-    props.push(buildObjectPropery('args', t.arrayExpression(schema.args.map(
-      arg => t.objectExpression(buildObjectPropsAst(arg))
-    ))))
+function inferArgType (e: t.Expression, getCode: GetCodeFn): TypeDescriptor {
+  if (AST_JSTYPE_MAP[e.type]) {
+    return {
+      type: AST_JSTYPE_MAP[e.type]
+    }
   }
-
-  if (schema.tags) {
-    props.push(buildObjectPropery('tags', t.arrayExpression(schema.tags.map(
-      tag => t.stringLiteral(tag)
-    ))))
+  if (e.type === 'AssignmentExpression') {
+    return inferArgType(e.right, getCode)
+  }
+  if (e.type === 'NewExpression' && e.callee.type === 'Identifier') {
+    return {
+      type: e.callee.name as JSType
+    }
   }
+  if (e.type === 'ArrayExpression' || e.type === 'TupleExpression') {
+    const itemTypes = e.elements
+      .filter(el => t.isExpression(el))
+      .flatMap(el => inferArgType(el as any, getCode).type)
+    return {
+      type: 'array',
+      items: { type: normalizeTypes(itemTypes) }
+    }
+  }
+  return {}
+}
 
-  return props
+function inferAnnotationType (ann: t.Identifier['typeAnnotation'], getCode: GetCodeFn): TypeDescriptor | null {
+  if (ann.type !== 'TSTypeAnnotation') { return null }
+  return inferTSType(ann.typeAnnotation, getCode)
 }
 
-function schemaToAst (schema) {
-  return t.objectExpression(schemaToPropsAst(schema))
+function inferTSType (tsType: t.TSType, getCode: GetCodeFn): TypeDescriptor | null {
+  if (tsType.type === 'TSParenthesizedType') {
+    return inferTSType(tsType.typeAnnotation, getCode)
+  }
+  if (tsType.type === 'TSTypeReference') {
+    if ('name' in tsType.typeName && tsType.typeName.name === 'Array') {
+      return {
+        type: 'array',
+        items: inferTSType(tsType.typeParameters.params[0], getCode)
+      }
+    }
+    return {
+      type: getCode(tsType.loc) as JSType
+    }
+  }
+  if (tsType.type === 'TSUnionType') {
+    return mergedTypes(...tsType.types.map(t => inferTSType(t, getCode)))
+  }
+  if (tsType.type === 'TSArrayType') {
+    return {
+      type: 'array',
+      items: inferTSType(tsType.elementType, getCode)
+    }
+  }
+  // if (tsType.type.endsWith('Keyword')) {
+  return {
+    type: getCode(tsType.loc) as JSType
+  }
+  // }
+  // return null
 }

src/types.ts

@@ -23,23 +23,28 @@ export type JSType =
 // eslint-disable-next-line no-use-before-define
 export type ResolveFn = ((value: any, get: (key: string) => any) => JSValue)
 
-export interface Schema {
-  id?: string,
+export interface TypeDescriptor {
   type?: JSType | JSType[]
-  items?: Schema
+  items?: TypeDescriptor | TypeDescriptor[]
+}
+
+export interface FunctionArg extends TypeDescriptor {
+  name?: string
+  default?: JSValue
+  optional?: boolean
+}
+
+export interface Schema extends TypeDescriptor {
+  id?: string,
   default?: JSValue
   resolve?: ResolveFn
   properties?: { [key: string]: Schema }
   title?: string
   description?: string
   $schema?: string
   tags?: string[]
-  args?: {
-    name: string
-    type?: string,
-    default?: any
-    optional?: boolean
-  }[]
+  args?: FunctionArg[]
+  returns?: TypeDescriptor,
 }
 
 export interface InputObject {