cat_wallet mypy fixes

chia/wallet/cat_wallet/cat_outer_puzzle.py

@@ -7,14 +7,14 @@
 from chia.types.coin_spend import CoinSpend
 from chia.util.ints import uint64
 from chia.wallet.cat_wallet.cat_utils import (
-    CAT_MOD,
     SpendableCAT,
     construct_cat_puzzle,
     match_cat_puzzle,
     unsigned_spend_bundle_for_spendable_cats,
 )
 from chia.wallet.lineage_proof import LineageProof
 from chia.wallet.puzzle_drivers import PuzzleInfo, Solver
+from chia.wallet.puzzles.cat_loader import CAT_MOD
 
 
 @dataclass(frozen=True)

chia/wallet/cat_wallet/cat_utils.py

@@ -36,15 +36,16 @@ def match_cat_puzzle(mod: Program, curried_args: Program) -> Optional[Iterator[P
     if it is, return the curried arguments
     """
     if mod == CAT_MOD:
-        return curried_args.as_iter()
+        ret: Iterator[Program] = curried_args.as_iter()
+        return ret
     else:
         return None
 
 
 def get_innerpuzzle_from_puzzle(puzzle: Program) -> Program:
     mod, curried_args = puzzle.uncurry()
     if mod == CAT_MOD:
-        return curried_args.rest().rest().first()
+        return curried_args.at("rrf")
     else:
         raise ValueError("Not a CAT puzzle")
 
@@ -60,7 +61,7 @@ def construct_cat_puzzle(
     return mod_code.curry(mod_code_hash, limitations_program_hash, inner_puzzle)
 
 
-def subtotals_for_deltas(deltas) -> List[int]:
+def subtotals_for_deltas(deltas: List[int]) -> List[int]:
     """
     Given a list of deltas corresponding to input coins, create the "subtotals" list
     needed in solutions spending those coins.
@@ -82,7 +83,9 @@ def subtotals_for_deltas(deltas) -> List[int]:
 def next_info_for_spendable_cat(spendable_cat: SpendableCAT) -> Program:
     c = spendable_cat.coin
     list = [c.parent_coin_info, spendable_cat.inner_puzzle.get_tree_hash(), c.amount]
-    return Program.to(list)
+    # ignoring hint error here for:
+    # https://github.com/Chia-Network/clvm/pull/102
+    return Program.to(list)  # type: ignore[no-any-return]
 
 
 # This should probably return UnsignedSpendBundle if that type ever exists
@@ -95,7 +98,7 @@ def unsigned_spend_bundle_for_spendable_cats(mod_code: Program, spendable_cat_li
     N = len(spendable_cat_list)
 
     # figure out what the deltas are by running the inner puzzles & solutions
-    deltas = []
+    deltas: List[int] = []
     for spend_info in spendable_cat_list:
         error, conditions, cost = conditions_dict_for_solution(
             spend_info.inner_puzzle, spend_info.inner_solution, INFINITE_COST

chia/wallet/cat_wallet/cat_wallet.py

@@ -5,7 +5,7 @@
 import time
 import traceback
 from secrets import token_bytes
-from typing import Any, Dict, List, Optional, Set, Tuple
+from typing import Any, Dict, List, Optional, Set, Tuple, TYPE_CHECKING
 
 from blspy import AugSchemeMPL, G2Element, G1Element
 
@@ -28,12 +28,12 @@
 from chia.wallet.cat_wallet.cat_constants import DEFAULT_CATS
 from chia.wallet.cat_wallet.cat_info import CATInfo, LegacyCATInfo
 from chia.wallet.cat_wallet.cat_utils import (
-    CAT_MOD,
     SpendableCAT,
     construct_cat_puzzle,
     match_cat_puzzle,
     unsigned_spend_bundle_for_spendable_cats,
 )
+from chia.wallet.puzzles.cat_loader import CAT_MOD
 from chia.wallet.cat_wallet.lineage_store import CATLineageStore
 from chia.wallet.coin_selection import select_coins
 from chia.wallet.derivation_record import DerivationRecord
@@ -54,11 +54,14 @@
 from chia.wallet.wallet_coin_record import WalletCoinRecord
 from chia.wallet.wallet_info import WalletInfo
 
+if TYPE_CHECKING:
+    from chia.wallet.wallet_state_manager import WalletStateManager
+
 # This should probably not live in this file but it's for experimental right now
 
 
 class CATWallet:
-    wallet_state_manager: Any
+    wallet_state_manager: WalletStateManager
     log: logging.Logger
     wallet_info: WalletInfo
     cat_info: CATInfo
@@ -72,8 +75,12 @@ def default_wallet_name_for_unknown_cat(limitations_program_hash_hex: str) -> st
 
     @staticmethod
     async def create_new_cat_wallet(
-        wallet_state_manager: Any, wallet: Wallet, cat_tail_info: Dict[str, Any], amount: uint64, name=None
-    ):
+        wallet_state_manager: WalletStateManager,
+        wallet: Wallet,
+        cat_tail_info: Dict[str, Any],
+        amount: uint64,
+        name: Optional[str] = None,
+    ) -> "CATWallet":
         self = CATWallet()
         self.cost_of_single_tx = None
         self.standard_wallet = wallet
@@ -162,10 +169,10 @@ async def create_new_cat_wallet(
 
     @staticmethod
     async def create_wallet_for_cat(
-        wallet_state_manager: Any,
+        wallet_state_manager: WalletStateManager,
         wallet: Wallet,
         limitations_program_hash_hex: str,
-        name=None,
+        name: Optional[str] = None,
     ) -> CATWallet:
         self = CATWallet()
         self.cost_of_single_tx = None
@@ -199,10 +206,10 @@ async def create_wallet_for_cat(
     @classmethod
     async def create_from_puzzle_info(
         cls,
-        wallet_state_manager: Any,
+        wallet_state_manager: WalletStateManager,
         wallet: Wallet,
         puzzle_driver: PuzzleInfo,
-        name=None,
+        name: Optional[str] = None,
     ) -> CATWallet:
         return await cls.create_wallet_for_cat(
             wallet_state_manager,
@@ -213,7 +220,7 @@ async def create_from_puzzle_info(
 
     @staticmethod
     async def create(
-        wallet_state_manager: Any,
+        wallet_state_manager: WalletStateManager,
         wallet: Wallet,
         wallet_info: WalletInfo,
     ) -> CATWallet:
@@ -259,19 +266,20 @@ async def get_confirmed_balance(self, record_list: Optional[Set[WalletCoinRecord
         self.log.info(f"Confirmed balance for cat wallet {self.id()} is {amount}")
         return uint64(amount)
 
-    async def get_unconfirmed_balance(self, unspent_records=None) -> uint128:
+    async def get_unconfirmed_balance(self, unspent_records: Optional[Set[WalletCoinRecord]] = None) -> uint128:
         return await self.wallet_state_manager.get_unconfirmed_balance(self.id(), unspent_records)
 
-    async def get_max_send_amount(self, records=None):
+    async def get_max_send_amount(self, records: Optional[Set[WalletCoinRecord]] = None) -> int:
         spendable: List[WalletCoinRecord] = list(await self.get_cat_spendable_coins())
         if len(spendable) == 0:
             return 0
         spendable.sort(reverse=True, key=lambda record: record.coin.amount)
         if self.cost_of_single_tx is None:
             coin = spendable[0].coin
             txs = await self.generate_signed_transaction(
-                [coin.amount], [coin.puzzle_hash], coins={coin}, ignore_max_send_amount=True
+                [uint64(coin.amount)], [coin.puzzle_hash], coins={coin}, ignore_max_send_amount=True
             )
+            assert txs[0].spend_bundle
             program: BlockGenerator = simple_solution_generator(txs[0].spend_bundle)
             # npc contains names of the coins removed, puzzle_hashes and their spend conditions
             result: NPCResult = get_name_puzzle_conditions(
@@ -297,18 +305,18 @@ async def get_max_send_amount(self, records=None):
 
         return total_amount
 
-    async def get_name(self):
+    async def get_name(self) -> str:
         return self.wallet_info.name
 
-    async def set_name(self, new_name: str):
+    async def set_name(self, new_name: str) -> None:
         new_info = dataclasses.replace(self.wallet_info, name=new_name)
         self.wallet_info = new_info
         await self.wallet_state_manager.user_store.update_wallet(self.wallet_info)
 
     def get_asset_id(self) -> str:
         return bytes(self.cat_info.limitations_program_hash).hex()
 
-    async def set_tail_program(self, tail_program: str):
+    async def set_tail_program(self, tail_program: str) -> None:
         assert Program.fromhex(tail_program).get_tree_hash() == self.cat_info.limitations_program_hash
         await self.save_info(
             CATInfo(
@@ -317,7 +325,7 @@ async def set_tail_program(self, tail_program: str):
             )
         )
 
-    async def coin_added(self, coin: Coin, height: uint32, peer: WSChiaConnection):
+    async def coin_added(self, coin: Coin, height: uint32, peer: WSChiaConnection) -> None:
         """Notification from wallet state manager that wallet has been received."""
         self.log.info(f"CAT wallet has been notified that {coin} was added")
 
@@ -340,7 +348,7 @@ async def coin_added(self, coin: Coin, height: uint32, peer: WSChiaConnection):
             except Exception as e:
                 self.log.debug(f"Exception: {e}, traceback: {traceback.format_exc()}")
 
-    async def puzzle_solution_received(self, coin_spend: CoinSpend, parent_coin: Coin):
+    async def puzzle_solution_received(self, coin_spend: CoinSpend, parent_coin: Coin) -> None:
         coin_name = coin_spend.coin.name()
         puzzle: Program = Program.from_bytes(bytes(coin_spend.puzzle_reveal))
         args = match_cat_puzzle(*puzzle.uncurry())
@@ -378,10 +386,10 @@ def puzzle_for_pk(self, pubkey: G1Element) -> Program:
         cat_puzzle: Program = construct_cat_puzzle(CAT_MOD, self.cat_info.limitations_program_hash, inner_puzzle)
         return cat_puzzle
 
-    async def get_new_cat_puzzle_hash(self):
+    async def get_new_cat_puzzle_hash(self) -> bytes32:
         return (await self.wallet_state_manager.get_unused_derivation_record(self.id())).puzzle_hash
 
-    async def get_spendable_balance(self, records=None) -> uint128:
+    async def get_spendable_balance(self, records: Optional[Set[WalletCoinRecord]] = None) -> uint128:
         coins = await self.get_cat_spendable_coins(records)
         amount = 0
         for record in coins:
@@ -410,7 +418,7 @@ async def get_pending_change_balance(self) -> uint64:
 
         return uint64(addition_amount)
 
-    async def get_cat_spendable_coins(self, records=None) -> List[WalletCoinRecord]:
+    async def get_cat_spendable_coins(self, records: Optional[Set[WalletCoinRecord]] = None) -> List[WalletCoinRecord]:
         result: List[WalletCoinRecord] = []
 
         record_list: Set[WalletCoinRecord] = await self.wallet_state_manager.get_spendable_coins_for_wallet(
@@ -508,7 +516,7 @@ async def convert_puzzle_hash(self, puzzle_hash: bytes32) -> bytes32:
         else:
             return (await self.inner_puzzle_for_cat_puzhash(puzzle_hash)).get_tree_hash()
 
-    async def get_lineage_proof_for_coin(self, coin) -> Optional[LineageProof]:
+    async def get_lineage_proof_for_coin(self, coin: Coin) -> Optional[LineageProof]:
         return await self.lineage_store.get_lineage_proof(coin.parent_coin_info)
 
     async def create_tandem_xch_tx(
@@ -567,7 +575,7 @@ async def generate_unsigned_spendbundle(
         payments: List[Payment],
         fee: uint64 = uint64(0),
         cat_discrepancy: Optional[Tuple[int, Program]] = None,  # (extra_delta, limitations_solution)
-        coins: Set[Coin] = None,
+        coins: Optional[Set[Coin]] = None,
         coin_announcements_to_consume: Optional[Set[Announcement]] = None,
         puzzle_announcements_to_consume: Optional[Set[Announcement]] = None,
         min_coin_amount: Optional[uint64] = None,
@@ -701,7 +709,7 @@ async def generate_signed_transaction(
         amounts: List[uint64],
         puzzle_hashes: List[bytes32],
         fee: uint64 = uint64(0),
-        coins: Set[Coin] = None,
+        coins: Optional[Set[Coin]] = None,
         ignore_max_send_amount: bool = False,
         memos: Optional[List[List[bytes]]] = None,
         coin_announcements_to_consume: Optional[Set[Announcement]] = None,
@@ -780,7 +788,7 @@ async def generate_signed_transaction(
 
         return tx_list
 
-    async def add_lineage(self, name: bytes32, lineage: Optional[LineageProof]):
+    async def add_lineage(self, name: bytes32, lineage: Optional[LineageProof]) -> None:
         """
         Lineage proofs are stored as a list of parent coins and the lineage proof you will need if they are the
         parent of the coin you are trying to spend. 'If I'm your parent, here's the info you need to spend yourself'
@@ -789,11 +797,11 @@ async def add_lineage(self, name: bytes32, lineage: Optional[LineageProof]):
         if lineage is not None:
             await self.lineage_store.add_lineage_proof(name, lineage)
 
-    async def remove_lineage(self, name: bytes32):
+    async def remove_lineage(self, name: bytes32) -> None:
         self.log.info(f"Removing parent {name} (probably had a non-CAT parent)")
         await self.lineage_store.remove_lineage_proof(name)
 
-    async def save_info(self, cat_info: CATInfo):
+    async def save_info(self, cat_info: CATInfo) -> None:
         self.cat_info = cat_info
         current_info = self.wallet_info
         data_str = bytes(cat_info).hex()

chia/wallet/cat_wallet/lineage_store.py

@@ -19,7 +19,7 @@ class CATLineageStore:
     table_name: str
 
     @classmethod
-    async def create(cls, db_wrapper: DBWrapper2, asset_id: str):
+    async def create(cls, db_wrapper: DBWrapper2, asset_id: str) -> "CATLineageStore":
         self = cls()
         self.table_name = f"lineage_proofs_{asset_id}"
         self.db_wrapper = db_wrapper
@@ -56,7 +56,8 @@ async def get_lineage_proof(self, coin_id: bytes32) -> Optional[LineageProof]:
             await cursor.close()
 
         if row is not None and row[0] is not None:
-            return LineageProof.from_bytes(row[1])
+            ret: LineageProof = LineageProof.from_bytes(row[1])
+            return ret
 
         return None
 

chia/wallet/puzzles/tails.py

@@ -9,11 +9,11 @@
 from chia.wallet.lineage_proof import LineageProof
 from chia.wallet.puzzles.load_clvm import load_clvm
 from chia.wallet.cat_wallet.cat_utils import (
-    CAT_MOD,
     construct_cat_puzzle,
     unsigned_spend_bundle_for_spendable_cats,
     SpendableCAT,
 )
+from chia.wallet.puzzles.cat_loader import CAT_MOD
 from chia.wallet.cat_wallet.cat_info import CATInfo
 from chia.wallet.transaction_record import TransactionRecord
 

chia/wallet/wallet_state_manager.py

@@ -1481,7 +1481,9 @@ async def add_new_wallet(self, wallet: Any, wallet_id: int, create_puzzle_hashes
             await self.create_more_puzzle_hashes()
         self.state_changed("wallet_created")
 
-    async def get_spendable_coins_for_wallet(self, wallet_id: int, records=None) -> Set[WalletCoinRecord]:
+    async def get_spendable_coins_for_wallet(
+        self, wallet_id: int, records: Optional[Set[WalletCoinRecord]] = None
+    ) -> Set[WalletCoinRecord]:
         if records is None:
             records = await self.coin_store.get_unspent_coins_for_wallet(wallet_id)