feat: infer function types (#1)

Co-authored-by: Daniel Roe <[email protected]>

playground/App.vue

@@ -34,7 +34,7 @@
       <!-- Output -->
       <div class="block">
         <div class="block-title">
-          <Tabs v-model="state.outputTab" :tabs="['loader', 'schema', 'docs', 'types', 'resolved']" />
+          <Tabs v-model="state.outputTab" :tabs="['loader', 'schema', 'types', 'docs', 'resolved']" />
           <span class="block-label">Output</span>
         </div>
         <!-- Schema -->
@@ -101,7 +101,7 @@ export default defineComponent({
   setup () {
     const state = persistedState({
       editorTab: 'reference',
-      outputTab: 'types',
+      outputTab: 'schema',
       ref: defaultReference,
       input: defaultInput
     })

playground/components/markdown.vue

@@ -28,6 +28,9 @@ export default defineComponent({
     const render = ref(true)
     const rendered = safeComputed(() => marked(ctx.value, {
       highlight (code, lang) {
+        if (lang === 'ts') {
+          lang = 'js'
+        }
         const _lang = prism.languages[lang]
         return _lang ? prism.highlight(code, _lang) : code
       }

playground/consts.ts

@@ -1,4 +1,6 @@
 export const defaultReference = `
+export function add (id: String, date = new Date(), append?: Boolean) {}
+
 export const config = {
     name: 'default',
     price: 12.5,

playground/utils.ts

@@ -1,13 +1,14 @@
 import { reactive, watch, computed } from 'vue'
 
 const globalKeys = Object.getOwnPropertyNames(globalThis)
-  .filter(key => key[0].toLocaleLowerCase() === key[0])
+  .filter(key => key[0].toLocaleLowerCase() === key[0] && key !== 'console')
 
 export function evaluateSource (src) {
   // Basic commonjs transform
   src = src
     .replace('export default', 'module.exports = ')
     .replace(/export (const|let|var) (\w+) ?= ?/g, 'exports.$2 = ')
+    .replace(/export (function|class) (\w+)/g, 'exports.$2 = $1 $2 ')
 
   // eslint-disable-next-line no-new-func
   const fn = Function(`

src/generator/dts.ts

@@ -13,7 +13,19 @@ const TYPE_MAP: Record<JSType, string> = {
   function: 'Function'
 }
 
-const SCHEMA_KEYS = ['items', 'default', 'resolve', 'properties', 'title', 'description', '$schema', 'type', 'id']
+const SCHEMA_KEYS = [
+  'items',
+  'default',
+  'resolve',
+  'properties',
+  'title',
+  'description',
+  '$schema',
+  'type',
+  'tags',
+  'args',
+  'id'
+]
 
 export function generateTypes (schema: Schema, name: string = 'Untyped') {
   return `interface ${name} {\n  ` + _genTypes(schema, ' ').join('\n ') + '\n}'
@@ -32,6 +44,8 @@ function _genTypes (schema: Schema, spaces: string): string[] {
       if (val.type === 'array') {
         const _type = getTsType(val.items.type)
         type = _type.includes('|') ? `(${_type})[]` : `${_type}[]`
+      } else if (val.type === 'function') {
+        type = genFunctionType(val)
       } else {
         type = getTsType(val.type)
       }
@@ -56,6 +70,24 @@ function getTsType (type: JSType | JSType[]): string {
   return (type && TYPE_MAP[type]) || 'any'
 }
 
+export function genFunctionType (schema: Schema) {
+  const args = schema.args.map((arg) => {
+    let argStr = arg.name
+    if (arg.optional) {
+      argStr += '?'
+    }
+    if (arg.type) {
+      argStr += ': ' + arg.type
+    }
+    if (arg.default) {
+      argStr += ' = ' + arg.default
+    }
+    return argStr
+  })
+
+  return `(${args.join(', ')}) => {}`
+}
+
 function generateJSDoc (schema: Schema): string[] {
   let buff = []
 
@@ -83,6 +115,12 @@ function generateJSDoc (schema: Schema): string[] {
     }
   }
 
+  if (Array.isArray(schema.tags)) {
+    for (const tag of schema.tags) {
+      buff.push('', tag)
+    }
+  }
+
   // Normalize new lines in values
   buff = buff.map(i => i.split('\n')).flat()
 

src/generator/md.ts

@@ -1,4 +1,5 @@
 import type { Schema } from '../types'
+import { genFunctionType } from './dts'
 
 export function generateMarkdown (schema: Schema) {
   return _generateMarkdown(schema, '', '').join('\n')
@@ -12,18 +13,29 @@ export function _generateMarkdown (schema: Schema, title: string, level) {
   if (schema.type === 'object') {
     for (const key in schema.properties) {
       const val = schema.properties[key] as Schema
-      lines.push(..._generateMarkdown(val, `\`${key}\``, level + '#'))
+      lines.push('', ..._generateMarkdown(val, `\`${key}\``, level + '#'))
     }
     return lines
   }
 
-  lines.push(`**Type**: \`${schema.type}\`    `)
+  // Type and default
+  lines.push(`- **Type**: \`${schema.type}\``)
   if ('default' in schema) {
-    lines.push(`**Default**: \`${JSON.stringify(schema.default)}\`    `)
+    lines.push(`- **Default**: \`${JSON.stringify(schema.default)}\``)
   }
+  lines.push('')
+
+  // Title
   if (schema.title) {
-    lines.push('', schema.title, '')
+    lines.push('> ' + schema.title, '')
+  }
+
+  // Signuture (function)
+  if (schema.type === 'function') {
+    lines.push('```ts', genFunctionType(schema), '```', '')
   }
+
+  // Description
   if (schema.description) {
     lines.push('', schema.description, '')
   }

src/loader/babel.ts

@@ -7,84 +7,161 @@ export default function babelPluginUntyped () {
     visitor: {
       ObjectProperty (p) {
         if (p.node.leadingComments) {
-          const { comments, blockTags } = parseJSDocs(
+          const schema = parseJSDocs(
             p.node.leadingComments
               .filter(c => c.type === 'CommentBlock')
               .map(c => c.value)
           )
 
-          const schema: Partial<Schema> = {}
-          if (comments.length) {
-            schema.title = comments.shift()
-          }
-          if (comments.length) {
-            schema.description = comments.join('\n')
-          }
-          for (const tag in blockTags) {
-            schema[tag] = blockTags[tag]
-          }
-
-          const schemaAstProps =
-            Object.entries(schema).map(e => t.objectProperty(t.identifier(e[0]), t.stringLiteral(e[1] as string)))
-          const schemaAst = t.objectExpression(schemaAstProps)
-
           if (p.node.value.type === 'ObjectExpression') {
             const schemaProp = p.node.value.properties.find(prop =>
               ('key' in prop) && prop.key.type === 'Identifier' && prop.key.name === '$schema'
             )
             if (schemaProp && 'value' in schemaProp) {
               if (schemaProp.value.type === 'ObjectExpression') {
                 // Object has $schema
-                schemaProp.value.properties.push(...schemaAstProps)
+                schemaProp.value.properties.push(...schemaToPropsAst(schema))
               } else {
                 // Object has $schema which is not an object
                 // SKIP
               }
             } else {
               // Object has not $schema
-              p.node.value.properties.unshift(t.objectProperty(t.identifier('$schema'), schemaAst))
+              p.node.value.properties.unshift(buildObjectPropery('$schema', schemaToAst(schema)))
             }
           } else {
             // Literal value
             p.node.value = t.objectExpression([
               t.objectProperty(t.identifier('$default'), p.node.value),
-              t.objectProperty(t.identifier('$schema'), schemaAst)
+              t.objectProperty(t.identifier('$schema'), schemaToAst(schema))
             ])
           }
           p.node.leadingComments = []
         }
+      },
+      FunctionDeclaration (p) {
+        // TODO: find associated jsdoc
+        const schema = parseJSDocs(
+          (p.node.leadingComments || [])
+            .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() || ''
+
+        // 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') {
+            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
+          }
+
+          if (param.type === 'AssignmentPattern') {
+            arg.default = getCode(param.right.loc)
+          }
+          schema.args.push(arg)
+        })
+
+        // 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)]))
       }
     }
   }
 }
 
-function parseJSDocs (input: string | string[]) {
+function parseJSDocs (input: string | string[]): Schema {
+  const schema: Schema = {}
+
   const lines = [].concat(input)
     .map(c => c.split('\n').map(l => l.replace(/^[\s*]+|[\s*]$/, '')))
     .flat()
     .filter(Boolean)
 
   const comments: string[] = []
-
   while (lines.length && !lines[0].startsWith('@')) {
     comments.push(lines.shift())
   }
+  if (comments.length === 1) {
+    schema.title = comments[0]
+  } else if (comments.length > 1) {
+    schema.title = comments[0]
+    schema.description = comments.splice(1).join('\n')
+  }
+
+  if (lines.length) {
+    schema.tags = []
+    for (const line of lines) {
+      schema.tags.push(line)
+    }
+  }
+
+  return schema
+}
+
+function valToAstLiteral (val: any) {
+  if (typeof val === 'string') {
+    return t.stringLiteral(val)
+  }
+  if (typeof val === 'boolean') {
+    return t.booleanLiteral(val)
+  }
+  if (typeof val === 'number') {
+    return t.numericLiteral(val)
+  }
+  return null
+}
 
-  const blockTags: Record<string, string> = {}
-  let lastTag = null
-  for (const line of lines) {
-    const m = line.match(/@(\w+) ?(.*)/)
-    if (m) {
-      blockTags[m[1]] = m[2]
-      lastTag = m[1]
-    } else {
-      blockTags[lastTag] =
-        (blockTags[lastTag] ? blockTags[lastTag] + '\n' : '') + line
+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))
     }
   }
+  return props
+}
+
+function buildObjectPropery (name, val) {
+  return t.objectProperty(t.identifier(name), val)
+}
 
-  return {
-    comments,
-    blockTags
+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))
+    ))))
   }
+
+  if (schema.tags) {
+    props.push(buildObjectPropery('tags', t.arrayExpression(schema.tags.map(
+      tag => t.stringLiteral(tag)
+    ))))
+  }
+
+  return props
+}
+
+function schemaToAst (schema) {
+  return t.objectExpression(schemaToPropsAst(schema))
 }

src/types.ts

@@ -33,6 +33,13 @@ export interface Schema {
   title?: string
   description?: string
   $schema?: string
+  tags?: string[]
+  args?: {
+    name: string
+    type?: string,
+    default?: any
+    optional?: boolean
+  }[]
 }
 
 export interface InputObject {