From a847edf22dfb6637f980f3863eefd1cd531e535e Mon Sep 17 00:00:00 2001
From: Jonas Daniels <jonas.daniels@outlook.com>
Date: Thu, 18 Jan 2024 15:12:36 +1300
Subject: [PATCH] feat: Add AddressType to AbiType Register

---
 packages/thirdweb/src/index.ts                            | 6 ++++++
 packages/thirdweb/src/react/wallet-provider.tsx           | 8 +++++---
 packages/thirdweb/src/rpc/methods.ts                      | 2 +-
 packages/thirdweb/src/transaction/actions/estimate-gas.ts | 5 +++--
 packages/thirdweb/src/transaction/actions/execute.ts      | 2 +-
 packages/thirdweb/src/wallets/interfaces/wallet.ts        | 2 +-
 packages/thirdweb/src/wallets/private-key.ts              | 3 ++-
 7 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/packages/thirdweb/src/index.ts b/packages/thirdweb/src/index.ts
index 0c43a0ee83d..abeb4d1fa3e 100644
--- a/packages/thirdweb/src/index.ts
+++ b/packages/thirdweb/src/index.ts
@@ -5,6 +5,12 @@ import { transaction, type TransactionOptions } from "./transaction/index.js";
 import { memoizePromise } from "./utils/promise.js";
 import type { AbiFunction } from "abitype";
 
+declare module "abitype" {
+  export interface Register {
+    AddressType: string;
+  }
+}
+
 export type CreateClientOptions =
   | {
       clientId: string;
diff --git a/packages/thirdweb/src/react/wallet-provider.tsx b/packages/thirdweb/src/react/wallet-provider.tsx
index fcef12b96c8..b1853ef5f4c 100644
--- a/packages/thirdweb/src/react/wallet-provider.tsx
+++ b/packages/thirdweb/src/react/wallet-provider.tsx
@@ -7,7 +7,7 @@ import {
 } from "react";
 import type { IWallet } from "../wallets/interfaces/wallet.js";
 
-export type WalletWithId = IWallet & { _id: string };
+export type WalletWithId = IWallet<any> & { _id: string };
 
 type WalletContext = {
   activeWallet: WalletWithId | null;
@@ -115,14 +115,14 @@ export function useConnect() {
   const [error, setError] = useState<Error | null>(null);
 
   const connect = useCallback(
-    async (options: IWallet | (() => Promise<IWallet>)) => {
+    async function <T>(options: IWallet<T> | (() => Promise<IWallet<T>>)) {
       // reset error state
       setError(null);
       if (typeof options !== "function") {
         const walletWithId = options as WalletWithId;
         walletWithId._id = fakeUuid();
         connectWallet(walletWithId);
-        return;
+        return walletWithId as IWallet<T>;
       }
 
       setIsConnecting(true);
@@ -132,11 +132,13 @@ export function useConnect() {
         const walletWithId = wallet as WalletWithId;
         walletWithId._id = fakeUuid();
         connectWallet(walletWithId);
+        return walletWithId as IWallet<T>;
       } catch (e) {
         setError(e as Error);
       } finally {
         setIsConnecting(false);
       }
+      return null;
     },
     [connectWallet],
   );
diff --git a/packages/thirdweb/src/rpc/methods.ts b/packages/thirdweb/src/rpc/methods.ts
index 5fb58917da6..36c83d3feed 100644
--- a/packages/thirdweb/src/rpc/methods.ts
+++ b/packages/thirdweb/src/rpc/methods.ts
@@ -1,6 +1,6 @@
+import type { Address } from "abitype";
 import type { RPCClient } from "./index.js";
 import {
-  type Address,
   type BlockNumber,
   type BlockTag,
   type RpcBlock,
diff --git a/packages/thirdweb/src/transaction/actions/estimate-gas.ts b/packages/thirdweb/src/transaction/actions/estimate-gas.ts
index 3ec9946f4e4..68ff29e1a7c 100644
--- a/packages/thirdweb/src/transaction/actions/estimate-gas.ts
+++ b/packages/thirdweb/src/transaction/actions/estimate-gas.ts
@@ -1,13 +1,14 @@
-import type { AbiFunction, Address } from "abitype";
+import type { AbiFunction } from "abitype";
 import type { Transaction } from "../index.js";
 import { getDefaultGasOverrides } from "../../gas/fee-data.js";
 import { encode } from "./encode.js";
 import { formatTransactionRequest, hexToBigInt } from "viem/utils";
 import type { IWallet } from "../../wallets/interfaces/wallet.js";
+import type { Address } from "viem";
 
 export async function estimateGas<
   const abiFn extends AbiFunction,
-  const wallet extends IWallet,
+  const wallet extends IWallet<any>,
 >(tx: Transaction<abiFn>, wallet?: wallet) {
   const rpcRequest = tx.client.rpc({ chainId: tx.inputs.chainId });
 
diff --git a/packages/thirdweb/src/transaction/actions/execute.ts b/packages/thirdweb/src/transaction/actions/execute.ts
index aa09cf16ec1..fd8e9e75640 100644
--- a/packages/thirdweb/src/transaction/actions/execute.ts
+++ b/packages/thirdweb/src/transaction/actions/execute.ts
@@ -4,7 +4,7 @@ import type { IWallet } from "../../wallets/interfaces/wallet.js";
 
 export async function execute<
   const abiFn extends AbiFunction,
-  const wallet extends IWallet,
+  const wallet extends IWallet<any>,
 >(tx: Transaction<abiFn>, wallet: wallet) {
   if (!wallet || !wallet.address) {
     throw new Error("no wallet to sign with");
diff --git a/packages/thirdweb/src/wallets/interfaces/wallet.ts b/packages/thirdweb/src/wallets/interfaces/wallet.ts
index 04863a43205..2d1d0705009 100644
--- a/packages/thirdweb/src/wallets/interfaces/wallet.ts
+++ b/packages/thirdweb/src/wallets/interfaces/wallet.ts
@@ -1,7 +1,7 @@
+import type { Address } from "abitype";
 import type { Transaction } from "../../transaction/index.js";
 import type {
   AbiFunction,
-  Address,
   Hex,
   SignableMessage,
   TransactionReceipt,
diff --git a/packages/thirdweb/src/wallets/private-key.ts b/packages/thirdweb/src/wallets/private-key.ts
index 85af77a9067..59b4811d33b 100644
--- a/packages/thirdweb/src/wallets/private-key.ts
+++ b/packages/thirdweb/src/wallets/private-key.ts
@@ -1,4 +1,5 @@
 import type {
+  Address,
   Hash,
   Hex,
   PrivateKeyAccount,
@@ -7,7 +8,7 @@ import type {
   TypedDataDefinition,
 } from "viem";
 import type { ThirdwebClient } from "../client/client.js";
-import type { AbiFunction, Address, TypedData } from "abitype";
+import type { AbiFunction, TypedData } from "abitype";
 import type { Transaction } from "../transaction/index.js";
 import type { IWallet } from "./interfaces/wallet.js";