Skip to content

Commit 20c64d0

Browse files
ElliotWood13ElliotWood13
committed
Refactor type casting and variable naming in cork_susde.py and cork_usde.py
- Updated type casting from Mapping to Dict for account balances. - Filtered AMM pools to include only those with term configurations. - Renamed variable lp_token_addr to lp_addr for consistency in both files.
1 parent 3a66662 commit 20c64d0

2 files changed

Lines changed: 18 additions & 16 deletions

File tree

‎integrations/cork_susde.py‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -602,7 +602,7 @@ def get_block_balances(
602602
if existing_block < block:
603603
prev_block = existing_block
604604
start = existing_block + 1
605-
account_bals = deepcopy(cast(Mapping[ChecksumAddress, Union[Decimal, float]], cache_copy_of_account_bals[prev_block]))
605+
account_bals = deepcopy(cast(Dict[ChecksumAddress, Union[Decimal, float]], cache_copy_of_account_bals[prev_block]))
606606
break
607607

608608
# Fetch pair config from self.start_block if not already done
@@ -710,6 +710,7 @@ def get_block_balances(
710710
# remaining reserves after end of epoch/term.
711711
amm_contract_function = self.amm_contract.functions.getReserves
712712
amm_pools = self.amm_balances_by_lp_token.values()
713+
amm_pools_with_tc = [p for p in amm_pools if p.term_config is not None]
713714
amm_calls = [
714715
(
715716
self.amm_contract,
@@ -719,7 +720,7 @@ def get_block_balances(
719720
amm_pool.term_config.share_token_addr,
720721
],
721722
)
722-
for amm_pool in amm_pools
723+
for amm_pool in amm_pools_with_tc
723724
]
724725
multicall_results = multicall_by_address(
725726
wb3=self.w3,
@@ -732,7 +733,7 @@ def get_block_balances(
732733
# - The `result[0]` is the total balance of the asset token in the AMM pool
733734
# - The `result[1]` is the total balance of the share token in the AMM pool
734735
for amm_pool, result in zip(
735-
amm_pools, multicall_results
736+
amm_pools_with_tc, multicall_results
736737
):
737738
# Update the total Ethena-asset and PSM-shares balance of each AMM pool
738739
amm_pool.total_assets = (result[0], result[1])
@@ -911,9 +912,9 @@ def get_block_balances(
911912
# attribute the Ethena asset balances to the respective LP token holders
912913
elif psm_pool.term_config is not None and account_addr == psm_pool.term_config.amm_pool_addr and amount:
913914
tc = psm_pool.term_config
914-
lp_token_addr = tc.amm_lp_token_addr
915-
if lp_token_addr is not None:
916-
amm_pool = self.amm_balances_by_lp_token[lp_token_addr]
915+
lp_addr = tc.amm_lp_token_addr
916+
if lp_addr is not None:
917+
amm_pool = self.amm_balances_by_lp_token[lp_addr]
917918
amm_ta = cast(Tuple[int, int], amm_pool.total_assets)
918919
# If there are other Uniswap V4 pools which manage PSM-shares,
919920
if account_shares > amm_ta[1]:
@@ -951,7 +952,7 @@ def get_block_balances(
951952
account_bals[account_addr] = Decimal(bal) + qty
952953
print(
953954
"LPT-LVT-holder:",
954-
lp_token_addr, vault_share_token_addr, account_addr,
955+
lp_addr, vault_share_token_addr, account_addr,
955956
"start:", bal,
956957
"in:", qty,
957958
"end:", account_bals[account_addr]
@@ -965,7 +966,7 @@ def get_block_balances(
965966
)
966967
account_bals[account_addr] = Decimal(bal) + qty
967968
print(
968-
"LPT-holder:", lp_token_addr, account_addr,
969+
"LPT-holder:", lp_addr, account_addr,
969970
"start:", bal,
970971
"in:", qty,
971972
"end:", account_bals[account_addr]

‎integrations/cork_usde.py‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -600,7 +600,7 @@ def get_block_balances(
600600
if existing_block < block:
601601
prev_block = existing_block
602602
start = existing_block + 1
603-
account_bals = deepcopy(cast(Mapping[ChecksumAddress, Union[Decimal, float]], cache_copy_of_account_bals[prev_block]))
603+
account_bals = deepcopy(cast(Dict[ChecksumAddress, Union[Decimal, float]], cache_copy_of_account_bals[prev_block]))
604604
break
605605

606606
# Fetch pair config from self.start_block if not already done
@@ -708,6 +708,7 @@ def get_block_balances(
708708
# remaining reserves after end of epoch/term.
709709
amm_contract_function = self.amm_contract.functions.getReserves
710710
amm_pools = self.amm_balances_by_lp_token.values()
711+
amm_pools_with_tc = [p for p in amm_pools if p.term_config is not None]
711712
amm_calls = [
712713
(
713714
self.amm_contract,
@@ -717,7 +718,7 @@ def get_block_balances(
717718
amm_pool.term_config.share_token_addr,
718719
],
719720
)
720-
for amm_pool in amm_pools
721+
for amm_pool in amm_pools_with_tc
721722
]
722723
multicall_results = multicall_by_address(
723724
wb3=self.w3,
@@ -730,7 +731,7 @@ def get_block_balances(
730731
# - The `result[0]` is the total balance of the asset token in the AMM pool
731732
# - The `result[1]` is the total balance of the share token in the AMM pool
732733
for amm_pool, result in zip(
733-
amm_pools, multicall_results
734+
amm_pools_with_tc, multicall_results
734735
):
735736
# Update the total Ethena-asset and PSM-shares balance of each AMM pool
736737
amm_pool.total_assets = (result[0], result[1])
@@ -909,9 +910,9 @@ def get_block_balances(
909910
# attribute the Ethena asset balances to the respective LP token holders
910911
elif psm_pool.term_config is not None and account_addr == psm_pool.term_config.amm_pool_addr and amount:
911912
tc = psm_pool.term_config
912-
lp_token_addr = tc.amm_lp_token_addr
913-
if lp_token_addr is not None:
914-
amm_pool = self.amm_balances_by_lp_token[lp_token_addr]
913+
lp_addr = tc.amm_lp_token_addr
914+
if lp_addr is not None:
915+
amm_pool = self.amm_balances_by_lp_token[lp_addr]
915916
amm_ta = cast(Tuple[int, int], amm_pool.total_assets)
916917
# If there are other Uniswap V4 pools which manage PSM-shares,
917918
# the total amount of PSM-shares at the UniV4 PoolManager address
@@ -952,7 +953,7 @@ def get_block_balances(
952953
account_bals[account_addr] = Decimal(bal) + qty
953954
print(
954955
"LPT-LVT-holder:",
955-
lp_token_addr, vault_share_token_addr, account_addr,
956+
lp_addr, vault_share_token_addr, account_addr,
956957
"start:", bal,
957958
"in:", qty,
958959
"end:", account_bals[account_addr]
@@ -966,7 +967,7 @@ def get_block_balances(
966967
)
967968
account_bals[account_addr] = Decimal(bal) + qty
968969
print(
969-
"LPT-holder:", lp_token_addr, account_addr,
970+
"LPT-holder:", lp_addr, account_addr,
970971
"start:", bal,
971972
"in:", qty,
972973
"end:", account_bals[account_addr]

0 commit comments

Comments
 (0)