fix: sor / sdk type issues

balancer · Oct 10, 2023 · 748395d · 748395d
1 parent 57f79a9
commit 748395d
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/balancer-js/src/modules/data/pool/onchain-data.ts b/balancer-js/src/modules/data/pool/onchain-data.ts
@@ -3,6 +3,7 @@ import { SubgraphPoolBase } from '@/.';
 import { Provider } from '@ethersproject/providers';
 import { formatFixed } from '@ethersproject/bignumber';
 import { SubgraphToken } from '@balancer-labs/sor';
+import { PoolToken, Pool } from '@/types';
 
 const abi = [
   'function getSwapFeePercentage() view returns (uint256)',
@@ -43,6 +44,11 @@ const getSwapFeeFn = (poolType: string) => {
   }
 };
 
+type GenericToken = SubgraphToken | PoolToken;
+type GenericPool = Omit<SubgraphPoolBase | Pool, 'tokens'> & {
+  tokens: GenericToken[];
+};
+
 interface OnchainData {
   poolTokens: [string[], string[]];
   totalShares: string;
@@ -133,7 +139,7 @@ const poolTypeCalls = (poolType: string, poolTypeVersion = 1) => {
   }
 };
 
-const merge = (pool: SubgraphPoolBase, result: OnchainData) => ({
+const merge = <T extends GenericPool>(pool: T, result: OnchainData) => ({
   ...pool,
   tokens: result.poolTokens
     ? pool.tokens.map((token) => {
@@ -209,16 +215,16 @@ export const fetchOnChainPoolData = async (
   return results;
 };
 
-export async function getOnChainBalances(
-  subgraphPoolsOriginal: SubgraphPoolBase[],
+export async function getOnChainBalances<T extends GenericPool>(
+  subgraphPoolsOriginal: T[],
   _multiAddress: string,
   vaultAddress: string,
   provider: Provider,
   batchSize = 1024
 ): Promise<T[]> {
   if (subgraphPoolsOriginal.length === 0) return subgraphPoolsOriginal;
 
-  const poolsWithOnchainData: SubgraphPoolBase[] = [];
+  const poolsWithOnchainData: T[] = [];
 
   const onchainData = (await fetchOnChainPoolData(
     subgraphPoolsOriginal,