@trpc/upgrade 0.0.0-alpha.2 → 0.0.0-alpha.26

package/README.md

@@ -2,6 +2,9 @@ run locally with source files
 
 ```sh
 DEV=1 pnpx tsx path/to/cli.ts
+
+# example
+cd examples/minimal-react/client && DEV=1 pnpx tsx ../../../packages/upgrade/src/bin/cli.ts --force --skipTanstackQuery --verbose
 ```
 
 or compiled

package/dist/cli.cjs

@@ -1,5 +1,5 @@
 #!/usr/bin/env node
-var path = require('path');
+var path = require('node:path');
 var cli$1 = require('@effect/cli');
 var platform = require('@effect/platform');
 var platformNode = require('@effect/platform-node');
@@ -10,22 +10,33 @@ function _interopDefault (e) { return e && e.__esModule ? e : { default: e }; }
 
 var path__default = /*#__PURE__*/_interopDefault(path);
 
-var version = "0.0.0-alpha.2";
+var version = "0.0.0-alpha.25";
 
-const assertCleanGitTree = platform.Command.string(platform.Command.make('git', 'status')).pipe(effect.Effect.filterOrFail(effect.String.includes('nothing to commit'), ()=>'Git tree is not clean, please commit your changes and try again, or run with `--force`'));
+const MakeCommand = (command, ...args)=>{
+    return platform.Command.make(command, ...args).pipe(platform.Command.workingDirectory(process.cwd()), platform.Command.runInShell(true));
+};
+const assertCleanGitTree = platform.Command.string(MakeCommand('git', 'status')).pipe().pipe(effect.Effect.filterOrFail(effect.String.includes('nothing to commit'), ()=>'Git tree is not clean, please commit your changes and try again, or run with `--force`'));
+const getPackageManager = ()=>effect.Match.value(process.env.npm_config_user_agent ?? 'npm').pipe(effect.Match.when(effect.String.startsWith('pnpm'), ()=>'pnpm'), effect.Match.when(effect.String.startsWith('yarn'), ()=>'yarn'), effect.Match.when(effect.String.startsWith('bun'), ()=>'bun'), effect.Match.orElse(()=>'npm'));
 const installPackage = (packageName)=>{
-    const packageManager = effect.Match.value(process.env.npm_config_user_agent ?? 'npm').pipe(effect.Match.when(effect.String.startsWith('pnpm'), ()=>'pnpm'), effect.Match.when(effect.String.startsWith('yarn'), ()=>'yarn'), effect.Match.when(effect.String.startsWith('bun'), ()=>'bun'), effect.Match.orElse(()=>'npm'));
-    return platform.Command.streamLines(platform.Command.make(packageManager, 'install', packageName)).pipe(effect.Stream.mapEffect(effect.Console.log), effect.Stream.runDrain);
+    const packageManager = getPackageManager();
+    return platform.Command.streamLines(MakeCommand(packageManager, 'install', packageName)).pipe(effect.Stream.mapEffect(effect.Console.log), effect.Stream.runDrain);
+};
+const uninstallPackage = (packageName)=>{
+    const packageManager = getPackageManager();
+    const uninstallCmd = packageManager === 'yarn' ? 'remove' : 'uninstall';
+    return platform.Command.streamLines(MakeCommand(packageManager, uninstallCmd, packageName)).pipe(effect.Stream.mapEffect(effect.Console.log), effect.Stream.runDrain);
 };
 const filterIgnored = (files)=>effect.Effect.gen(function*() {
-        const ignores = yield* platform.Command.string(platform.Command.make('git', 'check-ignore', '**/*')).pipe(effect.Effect.map((_)=>_.split('\n')));
-        yield* effect.Effect.log('All files in program:', files.map((_)=>_.fileName));
-        yield* effect.Effect.log('Ignored files:', ignores);
+        const ignores = yield* platform.Command.string(MakeCommand('git', 'check-ignore', '**/*').pipe(platform.Command.runInShell(true))).pipe(effect.Effect.tap((_)=>effect.Effect.logDebug('Ignored files output:', _)), effect.Effect.map((_)=>_.split('\n')));
+        yield* effect.Effect.logDebug('cwd:', process.cwd());
+        yield* effect.Effect.logDebug('All files in program:', files.map((_)=>_.fileName));
+        yield* effect.Effect.logDebug('Ignored files:', ignores);
         // Ignore "common files"
         const filteredSourcePaths = files.filter((source)=>source.fileName.startsWith(path__default.default.resolve()) && // only look ahead of current directory
             !source.fileName.includes('/trpc/packages/') && // relative paths when running codemod locally
+            !source.fileName.includes('/node_modules/') && // always ignore node_modules
             !ignores.includes(source.fileName)).map((source)=>source.fileName);
-        yield* effect.Effect.log('Filtered files:', filteredSourcePaths);
+        yield* effect.Effect.logDebug('Filtered files:', filteredSourcePaths);
         return filteredSourcePaths;
     });
 const TSProgram = effect.Effect.succeed(typescript.findConfigFile(process.cwd(), typescript.sys.fileExists)).pipe(effect.Effect.filterOrFail(effect.Predicate.isNotNullable, ()=>'No tsconfig found'), effect.Effect.tap((_)=>effect.Effect.logDebug('Using tsconfig', _)), effect.Effect.map((_)=>typescript.readConfigFile(_, typescript.sys.readFile)), effect.Effect.map((_)=>typescript.parseJsonConfigFileContent(_.config, typescript.sys, process.cwd())), effect.Effect.map((_)=>typescript.createProgram({
@@ -40,15 +51,22 @@ const force = cli$1.Options.boolean('force').pipe(cli$1.Options.withAlias('f'),
  * TODO: Instead of default values these should be detected automatically from the TS program
  */ const trpcFile = cli$1.Options.text('trpcFile').pipe(cli$1.Options.withAlias('f'), cli$1.Options.withDefault('~/trpc'), cli$1.Options.withDescription('Path to the trpc import file'));
 const trpcImportName = cli$1.Options.text('trpcImportName').pipe(cli$1.Options.withAlias('i'), cli$1.Options.withDefault('trpc'), cli$1.Options.withDescription('Name of the trpc import'));
+const skipTanstackQuery = cli$1.Options.boolean('skipTanstackQuery').pipe(cli$1.Options.withAlias('q'), cli$1.Options.withDefault(false), cli$1.Options.withDescription('Skip installing @trpc/tanstack-react-query package'));
+const verbose = cli$1.Options.boolean('verbose').pipe(cli$1.Options.withAlias('v'), cli$1.Options.withDefault(false), cli$1.Options.withDescription('Enable verbose logging'));
 const rootComamnd = cli$1.Command.make('upgrade', {
     force,
     trpcFile,
-    trpcImportName
+    trpcImportName,
+    skipTanstackQuery,
+    verbose
 }, (args)=>effect.Effect.gen(function*() {
+        if (args.verbose) {
+            yield* effect.Effect.log('Running upgrade with args:', args);
+        }
         if (!args.force) {
             yield* assertCleanGitTree;
         }
-        const transforms = yield* effect.Effect.map(cli$1.Prompt.multiSelect({
+        const transforms = yield* effect.pipe(cli$1.Prompt.multiSelect({
             message: 'Select transforms to run',
             choices: [
                 {
@@ -60,20 +78,26 @@ const rootComamnd = cli$1.Command.make('upgrade', {
                     value: require.resolve(transformPath('../transforms/provider.ts'))
                 }
             ]
-        }), // Make sure provider transform runs first if it's selected
-        effect.Array.sortWith((a)=>!a.includes('provider.ts'), effect.Order.boolean));
+        }), effect.Effect.flatMap((selected)=>{
+            if (selected.length === 0) {
+                return effect.Effect.fail(new Error('Please select at least one transform to run'));
+            }
+            return effect.Effect.succeed(selected);
+        }), effect.Effect.map(// Make sure provider transform runs first if it's selected
+        effect.Array.sortWith((a)=>!a.includes('provider.ts'), effect.Order.boolean)));
         const program = yield* TSProgram;
         const sourceFiles = program.getSourceFiles();
         const commitedFiles = yield* filterIgnored(sourceFiles);
         yield* effect.Effect.forEach(transforms, (transform)=>{
-            return effect.pipe(effect.Effect.log('Running transform', transform), effect.Effect.flatMap(()=>effect.Effect.tryPromise(async ()=>import('jscodeshift/src/Runner.js').then(({ run })=>run(transform, commitedFiles, {
-                            ...args,
-                            verbose: true
-                        })))), effect.Effect.map((_)=>effect.Effect.log('Transform result', _)));
+            return effect.pipe(effect.Effect.log('Running transform', transform), effect.Effect.flatMap(()=>effect.Effect.tryPromise(async ()=>import('jscodeshift/src/Runner.js').then(({ run })=>run(transform, commitedFiles, args)))), effect.Effect.map((_)=>effect.Effect.log('Transform result', _)));
         });
-        yield* effect.Effect.log('Installing @trpc/tanstack-react-query');
-        yield* installPackage('@trpc/tanstack-react-query');
-    }));
+        if (!args.skipTanstackQuery) {
+            yield* effect.Effect.log('Installing @trpc/tanstack-react-query');
+            yield* installPackage('@trpc/tanstack-react-query');
+            yield* effect.Effect.log('Uninstalling @trpc/react-query');
+            yield* uninstallPackage('@trpc/react-query');
+        }
+    }).pipe(effect.Logger.withMinimumLogLevel(args.verbose ? effect.LogLevel.Debug : effect.LogLevel.Info)));
 const cli = cli$1.Command.run(rootComamnd, {
     name: 'tRPC Upgrade CLI',
     version: `v${version}`

package/dist/transforms/hooksToOptions.cjs

@@ -1,6 +1,41 @@
 Object.defineProperty(exports, '__esModule', { value: true });
 
-/* eslint-disable @typescript-eslint/no-non-null-assertion */ /* eslint-disable no-console */ const hookToOptions = {
+/**
+ * Replaces the identifier for the root path key
+ * of a member expression
+ *
+ * For dot notation like `rootKey.x.y.z` the AST
+ * is constructed with the `rootKey` being nested deep
+ * inside a wrapper MemberExpression holding `rootKey.x`
+ * and so on
+ *
+ * This function helps replace the `rootKey` identifier with
+ * the provided identifier node
+ */ function replaceMemberExpressionRootIndentifier(j, expr, id) {
+    if (j.Identifier.check(expr.object)) {
+        expr.object = id;
+        return true;
+    }
+    if (j.MemberExpression.check(expr.object)) {
+        return replaceMemberExpressionRootIndentifier(j, expr.object, id);
+    }
+    return false;
+}
+
+/**
+ * Walks the path upwards to look for the closest parent
+ * of the mentioned type
+ */ function findParentOfType(path, type) {
+    if (!type.check(path.node)) {
+        return findParentOfType(path.parentPath, type);
+    }
+    if (!path.parent) {
+        return false;
+    }
+    return path;
+}
+
+const hookToOptions = {
     useQuery: {
         lib: '@tanstack/react-query',
         fn: 'queryOptions'
@@ -155,23 +190,36 @@ function transform(file, api, options) {
                     name: oldIdentifier.name
                 }).forEach((path)=>{
                     if (j.MemberExpression.check(path.parent?.parent?.node)) {
-                        const callExprPath = path.parent.parent.parent;
+                        const callExprPath = findParentOfType(path.parentPath, j.CallExpression);
+                        if (!callExprPath) {
+                            console.warn(`Failed to walk up the tree to find utilMethod call expression, on file: ${file.path}`, callExprPath, {
+                                start: path.node.loc?.start,
+                                end: path.node.loc?.end
+                            });
+                            return;
+                        }
                         const callExpr = callExprPath.node;
                         const memberExpr = callExpr.callee;
                         if (!j.CallExpression.check(callExpr) || !j.MemberExpression.check(memberExpr)) {
-                            console.warn('Failed to walk up the tree to find utilMethod call expression', callExpr);
+                            console.warn(`Failed to walk up the tree to find utilMethod with a \`trpc.PATH.<call>\`, on file: ${file.path}`, callExpr, {
+                                start: path.node.loc?.start,
+                                end: path.node.loc?.end
+                            });
                             return;
                         }
-                        if (!(j.MemberExpression.check(memberExpr.object) && j.Identifier.check(memberExpr.object.object) && j.Identifier.check(memberExpr.property) && memberExpr.property.name in utilMap)) {
+                        if (!(j.MemberExpression.check(memberExpr.object) && j.Identifier.check(memberExpr.property) && memberExpr.property.name in utilMap)) {
                             console.warn('Failed to identify utilMethod from proxy call expression', memberExpr);
                             return;
                         }
                         // Replace util.PATH.proxyMethod() with trpc.PATH.queryFilter()
                         const proxyMethod = memberExpr.property.name;
-                        memberExpr.object.object = j.identifier(trpcImportName);
+                        const replacedPath = replaceMemberExpressionRootIndentifier(j, memberExpr, j.identifier(trpcImportName));
+                        if (!replacedPath) {
+                            console.warn('Failed to wrap proxy call expression', memberExpr);
+                        }
                         memberExpr.property = j.identifier('queryFilter');
                         // Wrap it in queryClient.utilMethod()
-                        j(callExprPath).replaceWith(j.memberExpression(j.identifier('queryClient'), j.callExpression(j.identifier(utilMap[proxyMethod]), [
+                        callExprPath.replace(j.memberExpression(j.identifier('queryClient'), j.callExpression(j.identifier(utilMap[proxyMethod]), [
                             callExpr
                         ])));
                     }

package/dist/transforms/provider.cjs

@@ -2,6 +2,7 @@ Object.defineProperty(exports, '__esModule', { value: true });
 
 function transform(file, api, options) {
     const { trpcImportName } = options;
+    let routerName = undefined;
     const j = api.jscodeshift;
     const root = j(file.source);
     let dirtyFlag = false;
@@ -10,6 +11,8 @@ function transform(file, api, options) {
         const declaration = path.node.declarations[0];
         if (j.Identifier.check(declaration.id) && declaration.id.name === trpcImportName) {
             if (j.CallExpression.check(declaration.init) && j.Identifier.check(declaration.init.callee) && declaration.init.callee.name === 'createTRPCReact') {
+                // Get router name ( TODO : should probably get this from the TS compiler along with the import path)
+                routerName = declaration.init.original?.typeParameters?.params?.[0]?.typeName?.name;
                 // Replace the `createTRPCReact` call with `createTRPCContext`
                 declaration.init.callee.name = 'createTRPCContext';
                 // Destructure the result into `TRPCProvider` and `useTRPC`
@@ -46,7 +49,7 @@ function transform(file, api, options) {
             });
         });
     }
-    // Replace trpc.createClient with createTRPCClient
+    // Replace trpc.createClient with createTRPCClient<TRouter>
     root.find(j.CallExpression, {
         callee: {
             object: {
@@ -59,6 +62,11 @@ function transform(file, api, options) {
     }).forEach((path)=>{
         path.node.callee = j.identifier('createTRPCClient');
         dirtyFlag = true;
+        if (routerName) {
+            path.node.typeParameters = j.tsTypeParameterInstantiation([
+                j.tsTypeReference(j.identifier(routerName))
+            ]);
+        }
     });
     // Replace <trpc.Provider client={...} with <TRPCProvider trpcClient={...}
     root.find(j.JSXElement, {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@trpc/upgrade",
-  "version": "0.0.0-alpha.2",
+  "version": "0.0.0-alpha.26",
   "description": "Upgrade scripts for tRPC",
   "author": "juliusmarminge",
   "license": "MIT",
@@ -28,10 +28,10 @@
     "!**/__tests__"
   ],
   "dependencies": {
-    "@effect/cli": "0.48.25",
-    "@effect/platform": "0.69.25",
-    "@effect/platform-node": "0.64.27",
-    "effect": "3.10.16",
+    "@effect/cli": "0.55.0",
+    "@effect/platform": "0.76.0",
+    "@effect/platform-node": "0.72.0",
+    "effect": "3.12.11",
     "jscodeshift": "17.1.1",
     "typescript": "^5.6.2"
   },

package/src/bin/cli.ts

@@ -1,5 +1,5 @@
 /* eslint-disable @typescript-eslint/unbound-method */
-import path from 'path';
+import path from 'node:path';
 import { Command as CLICommand, Options, Prompt } from '@effect/cli';
 import { Command } from '@effect/platform';
 import { NodeContext, NodeRuntime } from '@effect/platform-node';
@@ -7,6 +7,8 @@ import {
   Array,
   Console,
   Effect,
+  Logger,
+  LogLevel,
   Match,
   Order,
   pipe,
@@ -24,39 +26,61 @@ import {
 } from 'typescript';
 import { version } from '../../package.json';
 
-const assertCleanGitTree = Command.string(Command.make('git', 'status')).pipe(
-  Effect.filterOrFail(
-    String.includes('nothing to commit'),
-    () =>
-      'Git tree is not clean, please commit your changes and try again, or run with `--force`',
-  ),
-);
+const MakeCommand = (command: string, ...args: string[]) => {
+  return Command.make(command, ...args).pipe(
+    Command.workingDirectory(process.cwd()),
+    Command.runInShell(true),
+  );
+};
 
-const installPackage = (packageName: string) => {
-  const packageManager = Match.value(
-    process.env.npm_config_user_agent ?? 'npm',
-  ).pipe(
+const assertCleanGitTree = Command.string(MakeCommand('git', 'status'))
+  .pipe()
+  .pipe(
+    Effect.filterOrFail(
+      String.includes('nothing to commit'),
+      () =>
+        'Git tree is not clean, please commit your changes and try again, or run with `--force`',
+    ),
+  );
+const getPackageManager = () =>
+  Match.value(process.env.npm_config_user_agent ?? 'npm').pipe(
     Match.when(String.startsWith('pnpm'), () => 'pnpm'),
     Match.when(String.startsWith('yarn'), () => 'yarn'),
     Match.when(String.startsWith('bun'), () => 'bun'),
     Match.orElse(() => 'npm'),
   );
+
+const installPackage = (packageName: string) => {
+  const packageManager = getPackageManager();
+  return Command.streamLines(
+    MakeCommand(packageManager, 'install', packageName),
+  ).pipe(Stream.mapEffect(Console.log), Stream.runDrain);
+};
+
+const uninstallPackage = (packageName: string) => {
+  const packageManager = getPackageManager();
+  const uninstallCmd = packageManager === 'yarn' ? 'remove' : 'uninstall';
   return Command.streamLines(
-    Command.make(packageManager, 'install', packageName),
+    MakeCommand(packageManager, uninstallCmd, packageName),
   ).pipe(Stream.mapEffect(Console.log), Stream.runDrain);
 };
 
 const filterIgnored = (files: readonly SourceFile[]) =>
   Effect.gen(function* () {
     const ignores = yield* Command.string(
-      Command.make('git', 'check-ignore', '**/*'),
-    ).pipe(Effect.map((_) => _.split('\n')));
+      MakeCommand('git', 'check-ignore', '**/*').pipe(Command.runInShell(true)),
+    ).pipe(
+      Effect.tap((_) => Effect.logDebug('Ignored files output:', _)),
+      Effect.map((_) => _.split('\n')),
+    );
 
-    yield* Effect.log(
+    yield* Effect.logDebug('cwd:', process.cwd());
+
+    yield* Effect.logDebug(
       'All files in program:',
       files.map((_) => _.fileName),
     );
-    yield* Effect.log('Ignored files:', ignores);
+    yield* Effect.logDebug('Ignored files:', ignores);
 
     // Ignore "common files"
     const filteredSourcePaths = files
@@ -64,11 +88,12 @@ const filterIgnored = (files: readonly SourceFile[]) =>
         (source) =>
           source.fileName.startsWith(path.resolve()) && // only look ahead of current directory
           !source.fileName.includes('/trpc/packages/') && // relative paths when running codemod locally
+          !source.fileName.includes('/node_modules/') && // always ignore node_modules
           !ignores.includes(source.fileName), // ignored files
       )
       .map((source) => source.fileName);
 
-    yield* Effect.log('Filtered files:', filteredSourcePaths);
+    yield* Effect.logDebug('Filtered files:', filteredSourcePaths);
 
     return filteredSourcePaths;
   });
@@ -114,20 +139,36 @@ const trpcImportName = Options.text('trpcImportName').pipe(
   Options.withDescription('Name of the trpc import'),
 );
 
+const skipTanstackQuery = Options.boolean('skipTanstackQuery').pipe(
+  Options.withAlias('q'),
+  Options.withDefault(false),
+  Options.withDescription('Skip installing @trpc/tanstack-react-query package'),
+);
+
+const verbose = Options.boolean('verbose').pipe(
+  Options.withAlias('v'),
+  Options.withDefault(false),
+  Options.withDescription('Enable verbose logging'),
+);
+
 const rootComamnd = CLICommand.make(
   'upgrade',
   {
     force,
     trpcFile,
     trpcImportName,
+    skipTanstackQuery,
+    verbose,
   },
   (args) =>
     Effect.gen(function* () {
+      if (args.verbose) {
+        yield* Effect.log('Running upgrade with args:', args);
+      }
       if (!args.force) {
         yield* assertCleanGitTree;
       }
-
-      const transforms = yield* Effect.map(
+      const transforms = yield* pipe(
         Prompt.multiSelect({
           message: 'Select transforms to run',
           choices: [
@@ -145,8 +186,18 @@ const rootComamnd = CLICommand.make(
             },
           ],
         }),
-        // Make sure provider transform runs first if it's selected
-        Array.sortWith((a) => !a.includes('provider.ts'), Order.boolean),
+        Effect.flatMap((selected) => {
+          if (selected.length === 0) {
+            return Effect.fail(
+              new Error('Please select at least one transform to run'),
+            );
+          }
+          return Effect.succeed(selected);
+        }),
+        Effect.map(
+          // Make sure provider transform runs first if it's selected
+          Array.sortWith((a) => !a.includes('provider.ts'), Order.boolean),
+        ),
       );
 
       const program = yield* TSProgram;
@@ -159,7 +210,7 @@ const rootComamnd = CLICommand.make(
           Effect.flatMap(() =>
             Effect.tryPromise(async () =>
               import('jscodeshift/src/Runner.js').then(({ run }) =>
-                run(transform, commitedFiles, { ...args, verbose: true }),
+                run(transform, commitedFiles, args),
               ),
             ),
           ),
@@ -167,9 +218,16 @@ const rootComamnd = CLICommand.make(
         );
       });
 
-      yield* Effect.log('Installing @trpc/tanstack-react-query');
-      yield* installPackage('@trpc/tanstack-react-query');
-    }),
+      if (!args.skipTanstackQuery) {
+        yield* Effect.log('Installing @trpc/tanstack-react-query');
+        yield* installPackage('@trpc/tanstack-react-query');
+
+        yield* Effect.log('Uninstalling @trpc/react-query');
+        yield* uninstallPackage('@trpc/react-query');
+      }
+    }).pipe(
+      Logger.withMinimumLogLevel(args.verbose ? LogLevel.Debug : LogLevel.Info),
+    ),
 );
 
 const cli = CLICommand.run(rootComamnd, {

package/src/lib/ast/modifiers.ts

@@ -0,0 +1,28 @@
+import type { Identifier, JSCodeshift, MemberExpression } from 'jscodeshift';
+
+/**
+ * Replaces the identifier for the root path key
+ * of a member expression
+ *
+ * For dot notation like `rootKey.x.y.z` the AST
+ * is constructed with the `rootKey` being nested deep
+ * inside a wrapper MemberExpression holding `rootKey.x`
+ * and so on
+ *
+ * This function helps replace the `rootKey` identifier with
+ * the provided identifier node
+ */
+export function replaceMemberExpressionRootIndentifier(
+  j: JSCodeshift,
+  expr: MemberExpression,
+  id: Identifier,
+) {
+  if (j.Identifier.check(expr.object)) {
+    expr.object = id;
+    return true;
+  }
+  if (j.MemberExpression.check(expr.object)) {
+    return replaceMemberExpressionRootIndentifier(j, expr.object, id);
+  }
+  return false;
+}

package/src/lib/ast/walkers.ts

@@ -0,0 +1,18 @@
+import type { ASTPath, JSCodeshift } from 'jscodeshift';
+
+/**
+ * Walks the path upwards to look for the closest parent
+ * of the mentioned type
+ */
+export function findParentOfType<TPath>(
+  path: ASTPath<unknown>,
+  type: JSCodeshift['AnyType'],
+): ASTPath<unknown> | false {
+  if (!type.check(path.node)) {
+    return findParentOfType(path.parentPath, type);
+  }
+  if (!path.parent) {
+    return false;
+  }
+  return path as ASTPath<TPath>;
+}

package/src/transforms/hooksToOptions.ts

@@ -11,6 +11,8 @@ import type {
   MemberExpression,
   Options,
 } from 'jscodeshift';
+import { replaceMemberExpressionRootIndentifier } from '../lib/ast/modifiers';
+import { findParentOfType } from '../lib/ast/walkers';
 
 interface TransformOptions extends Options {
   trpcImportName?: string;
@@ -205,7 +207,18 @@ export default function transform(
             .find(j.Identifier, { name: oldIdentifier.name })
             .forEach((path) => {
               if (j.MemberExpression.check(path.parent?.parent?.node)) {
-                const callExprPath = path.parent.parent.parent;
+                const callExprPath = findParentOfType<CallExpression>(
+                  path.parentPath,
+                  j.CallExpression,
+                );
+                if (!callExprPath) {
+                  console.warn(
+                    `Failed to walk up the tree to find utilMethod call expression, on file: ${file.path}`,
+                    callExprPath,
+                    { start: path.node.loc?.start, end: path.node.loc?.end },
+                  );
+                  return;
+                }
                 const callExpr = callExprPath.node as CallExpression;
                 const memberExpr = callExpr.callee as MemberExpression;
                 if (
@@ -213,8 +226,9 @@ export default function transform(
                   !j.MemberExpression.check(memberExpr)
                 ) {
                   console.warn(
-                    'Failed to walk up the tree to find utilMethod call expression',
+                    `Failed to walk up the tree to find utilMethod with a \`trpc.PATH.<call>\`, on file: ${file.path}`,
                     callExpr,
+                    { start: path.node.loc?.start, end: path.node.loc?.end },
                   );
                   return;
                 }
@@ -222,7 +236,6 @@ export default function transform(
                 if (
                   !(
                     j.MemberExpression.check(memberExpr.object) &&
-                    j.Identifier.check(memberExpr.object.object) &&
                     j.Identifier.check(memberExpr.property) &&
                     memberExpr.property.name in utilMap
                   )
@@ -236,11 +249,21 @@ export default function transform(
 
                 // Replace util.PATH.proxyMethod() with trpc.PATH.queryFilter()
                 const proxyMethod = memberExpr.property.name as ProxyMethod;
-                memberExpr.object.object = j.identifier(trpcImportName!);
+                const replacedPath = replaceMemberExpressionRootIndentifier(
+                  j,
+                  memberExpr,
+                  j.identifier(trpcImportName!),
+                );
+                if (!replacedPath) {
+                  console.warn(
+                    'Failed to wrap proxy call expression',
+                    memberExpr,
+                  );
+                }
                 memberExpr.property = j.identifier('queryFilter');
 
                 // Wrap it in queryClient.utilMethod()
-                j(callExprPath).replaceWith(
+                callExprPath.replace(
                   j.memberExpression(
                     j.identifier('queryClient'),
                     j.callExpression(j.identifier(utilMap[proxyMethod]), [

package/src/transforms/provider.ts

@@ -10,6 +10,7 @@ export default function transform(
   options: TransformOptions,
 ) {
   const { trpcImportName } = options;
+  let routerName: string | undefined = undefined;
 
   const j = api.jscodeshift;
   const root = j(file.source);
@@ -27,6 +28,11 @@ export default function transform(
         j.Identifier.check(declaration.init.callee) &&
         declaration.init.callee.name === 'createTRPCReact'
       ) {
+        // Get router name ( TODO : should probably get this from the TS compiler along with the import path)
+        routerName =
+          declaration.init.original?.typeParameters?.params?.[0]?.typeName
+            ?.name;
+
         // Replace the `createTRPCReact` call with `createTRPCContext`
         declaration.init.callee.name = 'createTRPCContext';
 
@@ -68,7 +74,7 @@ export default function transform(
       });
   }
 
-  // Replace trpc.createClient with createTRPCClient
+  // Replace trpc.createClient with createTRPCClient<TRouter>
   root
     .find(j.CallExpression, {
       callee: {
@@ -79,6 +85,12 @@ export default function transform(
     .forEach((path) => {
       path.node.callee = j.identifier('createTRPCClient');
       dirtyFlag = true;
+
+      if (routerName) {
+        (path.node as any).typeParameters = j.tsTypeParameterInstantiation([
+          j.tsTypeReference(j.identifier(routerName)),
+        ]);
+      }
     });
 
   // Replace <trpc.Provider client={...} with <TRPCProvider trpcClient={...}