Skip to content

Commit 741936e

Browse files
qnighyqnighy
committed
Improve function generation
1 parent 34b7208 commit 741936e

File tree

5 files changed

+155
-33
lines changed

5 files changed

+155
-33
lines changed

‎CHANGELOG.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
## Unreleased
22

3+
- Added
4+
- Add support for more type annotations on methods
5+
36
## 0.1.5
47

58
- Added

‎src/analysis/user_defined.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ export type UserDefinedFn = {
6262
type: "user_defined_function";
6363
localName?: string | undefined;
6464
init: FnInit;
65+
typeAnnotation?: NodePath<TSType> | undefined;
6566
sites: ThisFieldSite[];
6667
};
6768

@@ -160,6 +161,7 @@ export function analyzeUserDefined(
160161
fields.set(name, {
161162
type: "user_defined_function",
162163
init: fnInit,
164+
typeAnnotation: valInitType,
163165
sites: fieldSites,
164166
});
165167
} else if (isRefInit && !hasWrite) {

‎src/index.test.ts‎

Lines changed: 47 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -445,14 +445,24 @@ describe("react-declassify", () => {
445445
return null;
446446
}
447447
448-
foo = (x) => {
448+
foo = (x) => x + 42;
449+
bar = (x) => {
449450
return x + 42;
450-
}
451+
};
452+
baz = function(x) {
453+
return x + 42;
454+
};
451455
}
452456
`;
453457
const output = dedent`\
454458
const C = () => {
455-
function foo(x) {
459+
const foo = x => x + 42;
460+
461+
const bar = x => {
462+
return x + 42;
463+
};
464+
465+
function baz(x) {
456466
return x + 42;
457467
}
458468
@@ -604,6 +614,38 @@ describe("react-declassify", () => {
604614
`;
605615
expect(transform(input)).toBe(output);
606616
});
617+
618+
it("transforms method types", () => {
619+
const input = dedent`\
620+
class C extends React.Component {
621+
render() {
622+
return null;
623+
}
624+
625+
foo(x: number): number {
626+
return x + 42;
627+
}
628+
629+
bar: MyHandler = (x) => {
630+
return x + 42;
631+
}
632+
}
633+
`;
634+
const output = dedent`\
635+
const C: React.FC = () => {
636+
function foo(x: number): number {
637+
return x + 42;
638+
}
639+
640+
const bar: MyHandler = x => {
641+
return x + 42;
642+
};
643+
644+
return null;
645+
};
646+
`;
647+
expect(transform(input, { ts: true })).toBe(output);
648+
});
607649
});
608650

609651
describe("State transformation", () => {
@@ -731,9 +773,9 @@ describe("react-declassify", () => {
731773
const C = () => {
732774
const [foo, setFoo] = React.useState();
733775
734-
function reset() {
776+
const reset = () => {
735777
setFoo(42);
736-
}
778+
};
737779
};
738780
`;
739781
expect(transform(input)).toBe(output);

‎src/index.ts‎

Lines changed: 90 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
import type { ArrowFunctionExpression, Expression, Identifier, ImportDeclaration, MemberExpression, Pattern, RestElement, Statement, TSEntityName, TSType } from "@babel/types";
1+
import type { ArrowFunctionExpression, ClassMethod, ClassPrivateMethod, Expression, FunctionDeclaration, FunctionExpression, Identifier, ImportDeclaration, MemberExpression, ObjectMethod, Pattern, RestElement, Statement, TSEntityName, TSType, TSTypeAnnotation } from "@babel/types";
22
import type { NodePath, PluginObj, PluginPass } from "@babel/core";
3-
import { assignTypeAnnotation, assignTypeParameters, importName, isTS } from "./utils.js";
3+
import { assignReturnType, assignTypeAnnotation, assignTypeParameters, importName, isTS } from "./utils.js";
44
import { AnalysisError, analyzeBody, analyzeHead, ComponentBody, ComponentHead, needsProps, LibRef } from "./analysis.js";
55

66
type Options = {};
@@ -213,34 +213,28 @@ function transformClass(head: ComponentHead, body: ComponentBody, options: { ts:
213213
// Method definitions.
214214
if (field.init.type === "method") {
215215
const methNode = field.init.path.node;
216-
preamble.push(t.functionDeclaration(
217-
t.identifier(field.localName!),
218-
methNode.params as (Identifier | RestElement | Pattern)[],
219-
methNode.body,
220-
methNode.generator,
221-
methNode.async,
222-
));
216+
preamble.push(functionDeclarationFrom(babel, methNode, t.identifier(field.localName!)));
223217
} else {
224218
const methNode = field.init.initPath.node;
225-
if (methNode.type === "ArrowFunctionExpression") {
226-
preamble.push(t.functionDeclaration(
227-
t.identifier(field.localName!),
228-
methNode.params,
229-
methNode.body.type === "BlockStatement"
230-
? methNode.body
231-
: t.blockStatement([
232-
t.returnStatement(methNode.body)
233-
]),
234-
methNode.generator,
235-
methNode.async,
236-
));
219+
if (
220+
methNode.type === "FunctionExpression"
221+
&& !field.typeAnnotation
222+
) {
223+
preamble.push(functionDeclarationFrom(babel, methNode, t.identifier(field.localName!)));
237224
} else {
238-
preamble.push(t.functionDeclaration(
239-
t.identifier(field.localName!),
240-
methNode.params,
241-
methNode.body,
242-
methNode.generator,
243-
methNode.async,
225+
const expr =
226+
methNode.type === "FunctionExpression"
227+
? functionExpressionFrom(babel, methNode)
228+
: arrowFunctionExpressionFrom(babel, methNode);
229+
preamble.push(t.variableDeclaration(
230+
"const",
231+
[t.variableDeclarator(
232+
assignTypeAnnotation(
233+
t.identifier(field.localName!),
234+
field.typeAnnotation ? t.tsTypeAnnotation(field.typeAnnotation.node) : undefined
235+
),
236+
expr
237+
)]
244238
));
245239
}
246240
}
@@ -353,6 +347,75 @@ function getReactImport(
353347
return t.identifier(newName);
354348
}
355349

350+
type FunctionLike = FunctionDeclaration | FunctionExpression | ArrowFunctionExpression | ClassMethod | ClassPrivateMethod | ObjectMethod;
351+
352+
function functionName(node: FunctionLike): Identifier | undefined {
353+
switch (node.type) {
354+
case "FunctionDeclaration":
355+
case "FunctionExpression":
356+
return node.id ?? undefined;
357+
}
358+
}
359+
360+
function functionDeclarationFrom(
361+
babel: typeof import("@babel/core"),
362+
node: FunctionLike,
363+
name?: Identifier | null
364+
) {
365+
const { types: t } = babel;
366+
return assignReturnType(
367+
t.functionDeclaration(
368+
name ?? functionName(node),
369+
node.params as (Identifier | RestElement | Pattern)[],
370+
node.body.type === "BlockStatement"
371+
? node.body
372+
: t.blockStatement([
373+
t.returnStatement(node.body)
374+
]),
375+
node.generator,
376+
node.async,
377+
),
378+
node.returnType
379+
);
380+
}
381+
382+
function functionExpressionFrom(
383+
babel: typeof import("@babel/core"),
384+
node: FunctionLike,
385+
name?: Identifier | null
386+
) {
387+
const { types: t } = babel;
388+
return assignReturnType(
389+
t.functionExpression(
390+
name ?? functionName(node),
391+
node.params as (Identifier | RestElement | Pattern)[],
392+
node.body.type === "BlockStatement"
393+
? node.body
394+
: t.blockStatement([
395+
t.returnStatement(node.body)
396+
]),
397+
node.generator,
398+
node.async,
399+
),
400+
node.returnType
401+
);
402+
}
403+
404+
function arrowFunctionExpressionFrom(
405+
babel: typeof import("@babel/core"),
406+
node: FunctionLike
407+
) {
408+
const { types: t } = babel;
409+
return assignReturnType(
410+
t.arrowFunctionExpression(
411+
node.params as (Identifier | RestElement | Pattern)[],
412+
node.body,
413+
node.async,
414+
),
415+
node.returnType
416+
);
417+
}
418+
356419
/**
357420
* Refreshes recast's internal state to force generically printing comments.
358421
*/

‎src/utils.ts‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import type { NodePath, PluginPass } from "@babel/core";
2-
import type { ArrayPattern, AssignmentPattern, CallExpression, ClassAccessorProperty, ClassMethod, ClassPrivateMethod, ClassPrivateProperty, ClassProperty, Identifier, JSXOpeningElement, MemberExpression, NewExpression, ObjectMethod, ObjectPattern, ObjectProperty, OptionalCallExpression, RestElement, StaticBlock, StringLiteral, TaggedTemplateExpression, TSDeclareMethod, TSExpressionWithTypeArguments, TSImportType, TSInstantiationExpression, TSMethodSignature, TSPropertySignature, TSTypeAnnotation, TSTypeParameterInstantiation, TSTypeQuery, TSTypeReference } from "@babel/types";
2+
import type { ArrayPattern, ArrowFunctionExpression, AssignmentPattern, CallExpression, ClassAccessorProperty, ClassMethod, ClassPrivateMethod, ClassPrivateProperty, ClassProperty, FunctionDeclaration, FunctionExpression, Identifier, JSXOpeningElement, MemberExpression, NewExpression, Noop, ObjectMethod, ObjectPattern, ObjectProperty, OptionalCallExpression, RestElement, StaticBlock, StringLiteral, TaggedTemplateExpression, TSDeclareFunction, TSDeclareMethod, TSExpressionWithTypeArguments, TSImportType, TSInstantiationExpression, TSMethodSignature, TSPropertySignature, TSTypeAnnotation, TSTypeParameterInstantiation, TSTypeQuery, TSTypeReference, TypeAnnotation } from "@babel/types";
33

44
export function getOr<K, V>(m: Map<K, V>, k: K, getDefault: () => V): V {
55
if (m.has(k)) {
@@ -96,6 +96,18 @@ export function assignTypeAnnotation<T extends Annotatable>(
9696
});
9797
}
9898

99+
type ReturnTypeable =
100+
FunctionDeclaration | FunctionExpression | TSDeclareFunction | ArrowFunctionExpression | ObjectMethod | ClassMethod | ClassPrivateMethod | TSDeclareMethod;
101+
102+
export function assignReturnType<T extends ReturnTypeable>(
103+
node: T,
104+
returnType: TypeAnnotation | TSTypeAnnotation | Noop | null | undefined,
105+
): T {
106+
return Object.assign(node, {
107+
returnType,
108+
});
109+
}
110+
99111
type Paramable =
100112
CallExpression | NewExpression | TaggedTemplateExpression | OptionalCallExpression | JSXOpeningElement | TSTypeReference | TSTypeQuery | TSExpressionWithTypeArguments | TSInstantiationExpression | TSImportType;
101113

0 commit comments

Comments
 (0)