From 76cb65406c688ff784004c9ba908ebc647f8fac4 Mon Sep 17 00:00:00 2001 From: Daniel Roe Date: Sun, 28 Mar 2021 19:46:58 +0100 Subject: [PATCH] feat: improve function typing (#2) Co-authored-by: Pooya Parsa --- playground/consts.ts | 4 +- src/generator/dts.ts | 40 ++++++---- src/generator/md.ts | 4 +- src/loader/babel.ts | 174 +++++++++++++++++++++++++++++------------ src/types.ts | 23 +++--- src/utils.ts | 30 ++++++- test/transform.test.ts | 93 ++++++++++++++++++++++ test/types.test.ts | 31 +++++++- 8 files changed, 318 insertions(+), 81 deletions(-) create mode 100644 test/transform.test.ts diff --git a/playground/consts.ts b/playground/consts.ts index 7550542..8dc0896 100644 --- a/playground/consts.ts +++ b/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', diff --git a/src/generator/dts.ts b/src/generator/dts.ts index 4188814..e023bce 100644 --- a/src/generator/dts.ts +++ b/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 = { 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(', ') diff --git a/src/generator/md.ts b/src/generator/md.ts index 7b455d2..7558608 100644 --- a/src/generator/md.ts +++ b/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 diff --git a/src/loader/babel.ts b/src/loader/babel.ts index 464aa45..adf2acc 100644 --- a/src/loader/babel.ts +++ b/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 { 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> = { + 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 } diff --git a/src/types.ts b/src/types.ts index 0347408..b0cf04e 100644 --- a/src/types.ts +++ b/src/types.ts @@ -23,10 +23,19 @@ 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 } @@ -34,12 +43,8 @@ export interface Schema { description?: string $schema?: string tags?: string[] - args?: { - name: string - type?: string, - default?: any - optional?: boolean - }[] + args?: FunctionArg[] + returns?: TypeDescriptor, } export interface InputObject { diff --git a/src/utils.ts b/src/utils.ts index 68307f3..1f4a750 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -1,4 +1,4 @@ -import type { Schema, JSType } from './types' +import type { Schema, JSType, TypeDescriptor } from './types' export function escapeKey (val: string): string { return /^\w+$/.test(val) ? val : `"${val}"` @@ -76,3 +76,31 @@ export function getSchemaPath (schema: Schema, path) { } return schema } + +export function mergedTypes (...types: TypeDescriptor[]): TypeDescriptor { + types = types.filter(Boolean) + if (types.length === 0) { return {} } + if (types.length === 1) { return types[0] } + return { + type: normalizeTypes(types.map(t => t.type).flat().filter(Boolean)), + items: mergedTypes(...types.map(t => t.items).flat().filter(Boolean)) + } +} + +export function normalizeTypes (val: string[]) { + const arr = unique(val.filter(str => str)) + if (!arr.length || arr.includes('any')) { return undefined } + return (arr.length > 1) ? arr : arr[0] +} + +export function cachedFn (fn) { + let val + let resolved = false + return () => { + if (!resolved) { + val = fn() + resolved = true + } + return val + } +} diff --git a/test/transform.test.ts b/test/transform.test.ts new file mode 100644 index 0000000..3dcc508 --- /dev/null +++ b/test/transform.test.ts @@ -0,0 +1,93 @@ +import { transform } from '../src/loader/transform' + +describe('transform (functions)', () => { + it('creates correct types for simple function', () => { + const result = transform(` + export function add (id: string, date = new Date(), append?: boolean) {} + `) + + expectCodeToMatch(result, /export const add = ([\s\S]*)$/, { + $schema: { + type: 'function', + args: [{ + name: 'id', + type: 'string' + }, { + name: 'date', + type: 'Date' + }, { + name: 'append', + optional: true, + type: 'boolean' + }] + } + }) + }) + + it('infers correct types from defaults', () => { + const result = transform(` + export function add (test = ['42', 2], append?: false) {} + `) + + expectCodeToMatch(result, /export const add = ([\s\S]*)$/, { + $schema: { + type: 'function', + args: [{ + name: 'test', + type: 'array', + items: { + type: ['string', 'number'] + } + }, { + name: 'append', + type: 'false' + }] + } + }) + }) + + it('correctly uses a defined return type', () => { + const result = transform(` + export function add (): void {} + `) + + expectCodeToMatch(result, /export const add = ([\s\S]*)$/, { + $schema: { + type: 'function', + args: [], + returns: { + type: 'void' + } + } + }) + }) + + it('correctly handles a function assigned to a variable', () => { + const results = [transform(` + export const bob = function add (test: string): string {} + `), transform(` + export const bob = (test: string): string => {} + `)] + + results.forEach(result => expectCodeToMatch(result, /export const bob = ([\s\S]*)$/, { + $schema: { + type: 'function', + args: [{ + name: 'test', + type: 'string' + }], + returns: { + type: 'string' + } + } + })) + }) +}) + +function expectCodeToMatch (code: string, pattern: RegExp, expected: any) { + const [, result] = code.match(pattern) + expect(result).toBeDefined() + // eslint-disable-next-line + const obj = Function('"use strict";return (' + result.replace(/;$/, '') + ')')() + expect(obj).toMatchObject(expected) +} diff --git a/test/types.test.ts b/test/types.test.ts index 1b1d5c9..364d032 100644 --- a/test/types.test.ts +++ b/test/types.test.ts @@ -36,13 +36,13 @@ interface Untyped { expect(types).toBe(` interface Untyped { - empty: any[], + empty: Array, /** @default [1,2,3] */ - numbers: number[], + numbers: Array, /** @default [true,123] */ - mixed: (boolean | number)[], + mixed: Array, } `.trim()) }) @@ -53,4 +53,29 @@ interface Untyped { })) expect(types).toMatch('"*key": string') }) + + it('functions', () => { + const types = generateTypes(resolveSchema({ + add: { + $schema: { + type: 'function', + args: [{ + name: 'test', + type: 'Array', + optional: true + }, { + name: 'append', + type: 'boolean', + optional: true + }] + } + } + })) + + expect(types).toBe(` +interface Untyped { + add: (test?: Array, append?: boolean) => any, +} +`.trim()) + }) })