feat: improved chain types

packages/ensjs/src/clients/public.ts

@@ -1,13 +1,15 @@
 import {
   createClient,
-  type Chain,
   type Client,
   type ClientConfig,
   type PublicRpcSchema,
   type Transport,
 } from 'viem'
 import { addEnsContracts } from '../contracts/addEnsContracts.js'
-import type { ChainWithEns } from '../contracts/consts.js'
+import type {
+  ChainWithBaseContracts,
+  ChainWithEns,
+} from '../contracts/consts.js'
 import type { Prettify } from '../types.js'
 import { ensPublicActions, type EnsPublicActions } from './decorators/public.js'
 import {
@@ -17,7 +19,7 @@ import {
 
 export type EnsPublicClientConfig<
   TTransport extends Transport = Transport,
-  TChain extends Chain = Chain,
+  TChain extends ChainWithBaseContracts = ChainWithBaseContracts,
 > = Pick<
   ClientConfig<TTransport, TChain>,
   'batch' | 'key' | 'name' | 'pollingInterval' | 'transport'
@@ -56,7 +58,7 @@ export type EnsPublicClient<
  */
 export const createEnsPublicClient = <
   TTransport extends Transport,
-  TChain extends Chain,
+  TChain extends ChainWithBaseContracts,
 >({
   batch,
   chain,

packages/ensjs/src/clients/subgraph.ts

@@ -1,13 +1,15 @@
 import {
   createClient,
-  type Chain,
   type Client,
   type ClientConfig,
   type PublicRpcSchema,
   type Transport,
 } from 'viem'
 import { addEnsContracts } from '../contracts/addEnsContracts.js'
-import type { ChainWithEns } from '../contracts/consts.js'
+import type {
+  ChainWithBaseContracts,
+  ChainWithEns,
+} from '../contracts/consts.js'
 import type { Prettify } from '../types.js'
 import {
   ensSubgraphActions,
@@ -16,7 +18,7 @@ import {
 
 export type EnsSubgraphClientConfig<
   TTransport extends Transport = Transport,
-  TChain extends Chain = Chain,
+  TChain extends ChainWithBaseContracts = ChainWithBaseContracts,
 > = Pick<
   ClientConfig<TTransport, TChain>,
   'batch' | 'key' | 'name' | 'pollingInterval' | 'transport'
@@ -49,7 +51,7 @@ export type EnsSubgraphClient<
  */
 export const createEnsSubgraphClient = <
   TTransport extends Transport,
-  TChain extends Chain,
+  TChain extends ChainWithBaseContracts,
 >({
   batch,
   chain,

packages/ensjs/src/clients/wallet.ts

@@ -2,7 +2,6 @@ import {
   createWalletClient,
   type Account,
   type Address,
-  type Chain,
   type Client,
   type ClientConfig,
   type ParseAccount,
@@ -11,23 +10,30 @@ import {
   type WalletRpcSchema,
 } from 'viem'
 import { addEnsContracts } from '../contracts/addEnsContracts.js'
-import type { ChainWithEns } from '../contracts/consts.js'
-import type { Prettify } from '../types.js'
+import type {
+  ChainWithBaseContracts,
+  ChainWithEns,
+  CheckedChainWithEns,
+} from '../contracts/consts.js'
+import type { Assign, Prettify } from '../types.js'
 import { ensWalletActions, type EnsWalletActions } from './decorators/wallet.js'
 
 export type EnsWalletClientConfig<
   TTransport extends Transport,
-  TChain extends Chain,
+  TChain extends ChainWithBaseContracts,
   TAccountOrAddress extends Account | Address | undefined =
     | Account
     | Address
     | undefined,
-> = Pick<
-  ClientConfig<TTransport, TChain, TAccountOrAddress>,
-  'account' | 'chain' | 'key' | 'name' | 'pollingInterval' | 'transport'
-> & {
-  chain: TChain
-}
+> = Assign<
+  Pick<
+    ClientConfig<TTransport, TChain, TAccountOrAddress>,
+    'account' | 'chain' | 'key' | 'name' | 'pollingInterval' | 'transport'
+  >,
+  {
+    chain: TChain
+  }
+>
 
 export type EnsWalletClient<
   TTransport extends Transport = Transport,
@@ -61,7 +67,7 @@ export type EnsWalletClient<
  */
 export const createEnsWalletClient = <
   TTransport extends Transport,
-  TChain extends Chain,
+  TChain extends ChainWithBaseContracts,
   TAccountOrAddress extends Account | Address | undefined = undefined,
 >({
   account,
@@ -76,12 +82,12 @@ export const createEnsWalletClient = <
   TAccountOrAddress
 >): EnsWalletClient<
   TTransport,
-  ChainWithEns<TChain>,
+  CheckedChainWithEns<TChain>,
   ParseAccount<TAccountOrAddress>
 > => {
   return createWalletClient({
     account,
-    chain: addEnsContracts<TChain>(chain),
+    chain: addEnsContracts(chain),
     key,
     name,
     pollingInterval,

packages/ensjs/src/contracts/addEnsContracts.ts

@@ -4,7 +4,7 @@ import {
   addresses,
   subgraphs,
   supportedChains,
-  type ChainWithEns,
+  type CheckedChainWithEns,
   type SupportedChain,
 } from './consts.js'
 
@@ -22,7 +22,7 @@ import {
  *   transport: http(),
  * })
  */
-export const addEnsContracts = <TChain extends Chain>(chain: TChain) => {
+export const addEnsContracts = <const TChain extends Chain>(chain: TChain) => {
   if (!chain) throw new NoChainError()
   if (!supportedChains.includes(chain.network as SupportedChain))
     throw new UnsupportedNetworkError({
@@ -38,5 +38,5 @@ export const addEnsContracts = <TChain extends Chain>(chain: TChain) => {
     subgraphs: {
       ...subgraphs[chain.network as SupportedChain],
     },
-  } as ChainWithEns<TChain>
+  } as unknown as CheckedChainWithEns<TChain>
 }

packages/ensjs/src/contracts/consts.ts

@@ -6,6 +6,7 @@ import type {
   Transport,
   WalletClient,
 } from 'viem'
+import type { Assign, Prettify } from '../types.js'
 
 type ChainContract = {
   address: Address
@@ -177,12 +178,21 @@ export type ChainWithEns<TChain extends Chain = Chain> = TChain & {
   subgraphs: Subgraphs
 }
 
+export type ChainWithBaseContracts = Assign<
+  Omit<Chain, 'contracts'>,
+  {
+    contracts: BaseChainContracts
+  }
+>
+
 export type CheckedChainWithEns<TChain extends Chain> =
   TChain['network'] extends SupportedChain
-    ? TChain & {
-        contracts: (typeof addresses)[TChain['network']]
-        subgraphs: (typeof subgraphs)[TChain['network']]
-      }
+    ? TChain['contracts'] extends BaseChainContracts
+      ? TChain & {
+          contracts: Prettify<(typeof addresses)[TChain['network']]>
+          subgraphs: (typeof subgraphs)[TChain['network']]
+        }
+      : never
     : never
 
 export type ClientWithEns<

packages/ensjs/src/contracts/getChainContractAddress.ts

@@ -1,19 +1,29 @@
+import type { Chain } from 'viem'
 import { getChainContractAddress as _getChainContractAddress } from 'viem/utils'
-import type { ChainWithEns } from './consts.js'
 
-type ChainClient<TChain extends ChainWithEns> = {
-  chain: TChain
+type ExtractContract<TClient> = TClient extends {
+  chain: { contracts: infer C }
 }
+  ? C extends Record<string, { address: string }>
+    ? C
+    : never
+  : never
 
 export const getChainContractAddress = <
-  TChain extends ChainWithEns,
-  TClient extends ChainClient<TChain>,
+  const TClient extends { chain: Chain },
+  TContracts extends ExtractContract<TClient> = ExtractContract<TClient>,
+  TContract extends keyof TContracts = keyof TContracts,
 >({
   blockNumber,
   client,
   contract,
 }: {
   blockNumber?: bigint
   client: TClient
-  contract: Extract<keyof TChain['contracts'], string>
-}) => _getChainContractAddress({ blockNumber, chain: client.chain, contract })
+  contract: TContract
+}) =>
+  _getChainContractAddress({
+    blockNumber,
+    chain: client.chain,
+    contract: contract as string,
+  }) as TContracts[TContract]['address']

packages/ensjs/src/types.ts

@@ -11,6 +11,16 @@ export type Prettify<T> = {
   [K in keyof T]: T[K]
 } & {}
 
+type AssignI<T, U> = {
+  [K in keyof T as K extends keyof U
+    ? U[K] extends void
+      ? never
+      : K
+    : K]: K extends keyof U ? U[K] : T[K]
+}
+
+export type Assign<T, U> = AssignI<T, U> & U
+
 export type SimpleTransactionRequest = {
   [P in keyof Pick<TransactionRequest, 'to' | 'data'>]-?: Exclude<
     TransactionRequest[P],