(fix) take arrow function into account when adding return type (#1934)

#1932
also make ts a param

packages/language-server/package.json

@@ -57,7 +57,7 @@
         "prettier-plugin-svelte": "~2.9.0",
         "svelte": "^3.55.0",
         "svelte-preprocess": "~5.0.0",
-        "svelte2tsx": "~0.6.4",
+        "svelte2tsx": "~0.6.8",
         "typescript": "*",
         "vscode-css-languageservice": "~6.2.0",
         "vscode-html-languageservice": "~5.0.0",

packages/language-server/src/plugins/typescript/DocumentSnapshot.ts

@@ -537,6 +537,7 @@ export class JSOrTSDocumentSnapshot extends IdentityMapper implements DocumentSn
 
     private adjustText() {
         const result = internalHelpers.upsertKitFile(
+            ts,
             this.filePath,
             {
                 clientHooksPath: this.clientHooksPath,

packages/svelte2tsx/index.d.ts

@@ -128,14 +128,15 @@ export const internalHelpers: {
     )=> boolean,
     isParamsFile: (fileName: string, basename: string, paramsPath: string) =>boolean,
     upsertKitFile: (
+        _ts: typeof ts,
         fileName: string,
         kitFilesSettings: InternalHelpers.KitFilesSettings,
         getSource: () => ts.SourceFile | undefined,
         surround?: (code: string) => string
     ) => { text: string; addedCode: InternalHelpers.AddedCode[] } | undefined,
     toVirtualPos: (pos: number, addedCode: InternalHelpers.AddedCode[]) => number,
     toOriginalPos: (pos: number, addedCode: InternalHelpers.AddedCode[]) => {pos: number; inGenerated: boolean},
-    findExports: (source: ts.SourceFile, isTsFile: boolean) => Map<
+    findExports: (_ts: typeof ts, source: ts.SourceFile, isTsFile: boolean) => Map<
         string,
         | {
             type: 'function';

packages/svelte2tsx/package.json

@@ -1,6 +1,6 @@
 {
     "name": "svelte2tsx",
-    "version": "0.6.4",
+    "version": "0.6.8",
     "description": "Convert Svelte components to TSX for type checking",
     "author": "David Pershouse",
     "license": "MIT",

packages/svelte2tsx/src/helpers/sveltekit.ts

@@ -1,7 +1,9 @@
 import path from 'path';
-import ts from 'typescript';
+import type ts from 'typescript';
 import { findExports } from './typescript';
 
+type _ts = typeof ts;
+
 export interface AddedCode {
     generatedPos: number;
     originalPos: number;
@@ -87,29 +89,39 @@ export function isParamsFile(fileName: string, basename: string, paramsPath: str
 }
 
 export function upsertKitFile(
+    ts: _ts,
     fileName: string,
     kitFilesSettings: KitFilesSettings,
     getSource: () => ts.SourceFile | undefined,
     surround: (text: string) => string = (text) => text
 ): { text: string; addedCode: AddedCode[] } {
     let basename = path.basename(fileName);
     const result =
-        upserKitRouteFile(fileName, basename, getSource, surround) ??
+        upserKitRouteFile(ts, fileName, basename, getSource, surround) ??
         upserKitServerHooksFile(
+            ts,
             fileName,
             basename,
             kitFilesSettings.serverHooksPath,
             getSource,
             surround
         ) ??
         upserKitClientHooksFile(
+            ts,
             fileName,
             basename,
             kitFilesSettings.clientHooksPath,
             getSource,
             surround
         ) ??
-        upserKitParamsFile(fileName, basename, kitFilesSettings.paramsPath, getSource, surround);
+        upserKitParamsFile(
+            ts,
+            fileName,
+            basename,
+            kitFilesSettings.paramsPath,
+            getSource,
+            surround
+        );
     if (!result) {
         return;
     }
@@ -128,6 +140,7 @@ export function upsertKitFile(
 }
 
 function upserKitRouteFile(
+    ts: _ts,
     fileName: string,
     basename: string,
     getSource: () => ts.SourceFile | undefined,
@@ -144,7 +157,7 @@ function upserKitRouteFile(
     };
 
     const isTsFile = basename.endsWith('.ts');
-    const exports = findExports(source, isTsFile);
+    const exports = findExports(ts, source, isTsFile);
 
     // add type to load function if not explicitly typed
     const load = exports.get('load');
@@ -175,6 +188,7 @@ function upserKitRouteFile(
     // add types to GET/PUT/POST/PATCH/DELETE/OPTIONS if not explicitly typed
     const insertApiMethod = (name: string) => {
         addTypeToFunction(
+            ts,
             exports,
             surround,
             insert,
@@ -194,6 +208,7 @@ function upserKitRouteFile(
 }
 
 function upserKitParamsFile(
+    ts: _ts,
     fileName: string,
     basename: string,
     paramsPath: string,
@@ -213,14 +228,15 @@ function upserKitParamsFile(
     };
 
     const isTsFile = basename.endsWith('.ts');
-    const exports = findExports(source, isTsFile);
+    const exports = findExports(ts, source, isTsFile);
 
-    addTypeToFunction(exports, surround, insert, 'match', 'string', 'boolean');
+    addTypeToFunction(ts, exports, surround, insert, 'match', 'string', 'boolean');
 
     return { addedCode, originalText: source.getFullText() };
 }
 
 function upserKitClientHooksFile(
+    ts: _ts,
     fileName: string,
     basename: string,
     clientHooksPath: string,
@@ -240,9 +256,10 @@ function upserKitClientHooksFile(
     };
 
     const isTsFile = basename.endsWith('.ts');
-    const exports = findExports(source, isTsFile);
+    const exports = findExports(ts, source, isTsFile);
 
     addTypeToFunction(
+        ts,
         exports,
         surround,
         insert,
@@ -254,6 +271,7 @@ function upserKitClientHooksFile(
 }
 
 function upserKitServerHooksFile(
+    ts: _ts,
     fileName: string,
     basename: string,
     serverHooksPath: string,
@@ -273,10 +291,10 @@ function upserKitServerHooksFile(
     };
 
     const isTsFile = basename.endsWith('.ts');
-    const exports = findExports(source, isTsFile);
+    const exports = findExports(ts, source, isTsFile);
 
     const addType = (name: string, type: string) => {
-        addTypeToFunction(exports, surround, insert, name, type);
+        addTypeToFunction(ts, exports, surround, insert, name, type);
     };
 
     addType('handleError', `import('@sveltejs/kit').HandleServerError`);
@@ -310,6 +328,7 @@ function addTypeToVariable(
 }
 
 function addTypeToFunction(
+    ts: _ts,
     exports: Map<
         string,
         | {
@@ -331,9 +350,11 @@ function addTypeToFunction(
         const paramInsertion = surround(!returnType ? `: Parameters<${type}>[0]` : `: ${type}`);
         insert(paramPos, paramInsertion);
         if (!fn.node.type && fn.node.body) {
-            const returnPos = fn.node.body.getStart();
+            const returnPos = ts.isArrowFunction(fn.node)
+                ? fn.node.equalsGreaterThanToken.getStart()
+                : fn.node.body.getStart();
             const returnInsertion = surround(
-                !returnType ? `: ReturnType<${type}>` : `: ${returnType}`
+                !returnType ? `: ReturnType<${type}> ` : `: ${returnType} `
             );
             insert(returnPos, returnInsertion);
         }

packages/svelte2tsx/src/helpers/typescript.ts

@@ -1,9 +1,11 @@
-import ts from 'typescript';
+import type ts from 'typescript';
+
+type _ts = typeof ts;
 
 /**
  * Finds the top level const/let/function exports of a source file.
  */
-export function findExports(source: ts.SourceFile, isTsFile: boolean) {
+export function findExports(ts: _ts, source: ts.SourceFile, isTsFile: boolean) {
     const exports = new Map<
         string,
         | {
@@ -28,7 +30,7 @@ export function findExports(source: ts.SourceFile, isTsFile: boolean) {
             exports.set(statement.name.text, {
                 type: 'function',
                 node: statement,
-                hasTypeDefinition: hasTypedParameter(statement, isTsFile)
+                hasTypeDefinition: hasTypedParameter(ts, statement, isTsFile)
             });
         }
         if (
@@ -59,7 +61,7 @@ export function findExports(source: ts.SourceFile, isTsFile: boolean) {
                 exports.set(declaration.name.getText(), {
                     type: 'function',
                     node,
-                    hasTypeDefinition: hasTypeDefinition || hasTypedParameter(node, isTsFile)
+                    hasTypeDefinition: hasTypeDefinition || hasTypedParameter(ts, node, isTsFile)
                 });
             } else {
                 exports.set(declaration.name.getText(), {
@@ -75,6 +77,7 @@ export function findExports(source: ts.SourceFile, isTsFile: boolean) {
 }
 
 function hasTypedParameter(
+    ts: _ts,
     node: ts.FunctionDeclaration | ts.ArrowFunction | ts.FunctionExpression,
     isTsFile: boolean
 ): boolean {

packages/svelte2tsx/test/helpers/index.ts

@@ -0,0 +1,44 @@
+import * as assert from 'assert';
+import ts from 'typescript';
+import { internalHelpers } from '../../src';
+
+describe('Internal Helpers - upsertKitFile', () => {
+    function upsert(file: string, source: string, expected: string) {
+        const sourceFile = ts.createSourceFile('d', source, ts.ScriptTarget.Latest, true);
+        const result = internalHelpers.upsertKitFile(
+            ts,
+            file,
+            {
+                clientHooksPath: 'hooks.client',
+                paramsPath: 'params',
+                serverHooksPath: 'hooks.server'
+            },
+            () => sourceFile
+        );
+        assert.strictEqual(result?.text, expected);
+    }
+
+    it('upserts +page.ts function', () => {
+        upsert(
+            '+page.ts',
+            `export function load(e) { return e; }`,
+            `export function load(e: import('./$types').PageLoadEvent) { return e; }`
+        );
+    });
+
+    it('upserts handle hook const', () => {
+        upsert(
+            'hooks.server.ts',
+            `export const handle = async ({ event, resolve }) => {};`,
+            `export const handle = async ({ event, resolve }: Parameters<import('@sveltejs/kit').Handle>[0]) : ReturnType<import('@sveltejs/kit').Handle> => {};`
+        );
+    });
+
+    it('upserts GET async function', () => {
+        upsert(
+            '+server.ts',
+            `export async function GET(e) {}`,
+            `export async function GET(e: import('./$types').RequestEvent) : Response | Promise<Response> {}`
+        );
+    });
+});

packages/typescript-plugin/package.json

@@ -24,6 +24,6 @@
     },
     "dependencies": {
         "@jridgewell/sourcemap-codec": "^1.4.14",
-        "svelte2tsx": "~0.6.4"
+        "svelte2tsx": "~0.6.8"
     }
 }

packages/typescript-plugin/src/language-service/diagnostics.ts

@@ -170,7 +170,7 @@ function getKitDiagnostics<
             isKitRouteExportAllowedIn(basename, kitExports[key])
         );
         if (source && basename.startsWith('+')) {
-            const exports = internalHelpers.findExports(source, /* irrelevant */ false);
+            const exports = internalHelpers.findExports(ts, source, /* irrelevant */ false);
             for (const exportName of exports.keys()) {
                 if (!validExports.includes(exportName) && !exportName.startsWith('_')) {
                     const node = exports.get(exportName)!.node;

packages/typescript-plugin/src/language-service/sveltekit.ts

@@ -581,6 +581,7 @@ function getProxiedLanguageService(info: ts.server.PluginCreateInfo, ts: _ts, lo
 
         upsertKitFile(fileName: string) {
             const result = internalHelpers.upsertKitFile(
+                ts,
                 fileName,
                 {
                     clientHooksPath: this.clientHooksPath,