Puzzle driver type annotations (#13190)

* annotate CAT outer puzzle driver

* add type annotation to transfer_program_puzzle

* add type annotations to singleton_outer_puzzle

* add type annotations to ownership_outer_puzzle

* add type annotations to metadata_outer_puzzle

* remove unused _asset_id field in puzzle drivers

chia/wallet/cat_wallet/cat_outer_puzzle.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, List, Optional
+from typing import Callable, List, Optional
 
 from chia.types.blockchain_format.coin import Coin
 from chia.types.blockchain_format.program import Program
@@ -19,12 +19,11 @@ from chia.wallet.puzzles.cat_loader import CAT_MOD
 
 @dataclass(frozen=True)
 class CATOuterPuzzle:
-    _match: Any
-    _asset_id: Any
-    _construct: Any
-    _solve: Any
-    _get_inner_puzzle: Any
-    _get_inner_solution: Any
+    _match: Callable[[Program], Optional[PuzzleInfo]]
+    _construct: Callable[[PuzzleInfo, Program], Program]
+    _solve: Callable[[PuzzleInfo, Solver, Program, Program], Program]
+    _get_inner_puzzle: Callable[[PuzzleInfo, Program], Optional[Program]]
+    _get_inner_solution: Callable[[PuzzleInfo, Program], Optional[Program]]
 
     def match(self, puzzle: Program) -> Optional[PuzzleInfo]:
         args = match_cat_puzzle(*puzzle.uncurry())
@@ -45,16 +44,18 @@ class CATOuterPuzzle:
         if args is None:
             raise ValueError("This driver is not for the specified puzzle reveal")
         _, _, inner_puzzle = args
-        if constructor.also() is not None:
-            deep_inner_puzzle: Optional[Program] = self._get_inner_puzzle(constructor.also(), inner_puzzle)
+        also = constructor.also()
+        if also is not None:
+            deep_inner_puzzle: Optional[Program] = self._get_inner_puzzle(also, inner_puzzle)
             return deep_inner_puzzle
         else:
             return inner_puzzle
 
     def get_inner_solution(self, constructor: PuzzleInfo, solution: Program) -> Optional[Program]:
         my_inner_solution: Program = solution.first()
-        if constructor.also():
-            deep_inner_solution: Optional[Program] = self._get_inner_solution(constructor.also(), my_inner_solution)
+        also = constructor.also()
+        if also:
+            deep_inner_solution: Optional[Program] = self._get_inner_solution(also, my_inner_solution)
             return deep_inner_solution
         else:
             return my_inner_solution
@@ -63,8 +64,9 @@ class CATOuterPuzzle:
         return bytes32(constructor["tail"])
 
     def construct(self, constructor: PuzzleInfo, inner_puzzle: Program) -> Program:
-        if constructor.also() is not None:
-            inner_puzzle = self._construct(constructor.also(), inner_puzzle)
+        also = constructor.also()
+        if also is not None:
+            inner_puzzle = self._construct(also, inner_puzzle)
         return construct_cat_puzzle(CAT_MOD, constructor["tail"], inner_puzzle)
 
     def solve(self, constructor: PuzzleInfo, solver: Solver, inner_puzzle: Program, inner_solution: Program) -> Program:
@@ -91,9 +93,10 @@ class CATOuterPuzzle:
                 target_coin = coin
             parent_spend: CoinSpend = CoinSpend.from_bytes(spend_prog.as_python())
             parent_coin: Coin = parent_spend.coin
-            if constructor.also() is not None:
-                puzzle = self._construct(constructor.also(), puzzle)
-                solution = self._solve(constructor.also(), solver, inner_puzzle, inner_solution)
+            also = constructor.also()
+            if also is not None:
+                puzzle = self._construct(also, puzzle)
+                solution = self._solve(also, solver, inner_puzzle, inner_solution)
             args = match_cat_puzzle(*parent_spend.puzzle_reveal.to_program().uncurry())
             assert args is not None
             _, _, parent_inner_puzzle = args

chia/wallet/nft_wallet/metadata_outer_puzzle.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, List, Optional, Tuple
+from typing import Callable, List, Optional, Tuple
 
 from clvm_tools.binutils import disassemble
 
@@ -31,12 +31,11 @@ def solution_for_metadata_layer(amount: uint64, inner_solution: Program) -> Prog
 
 @dataclass(frozen=True)
 class MetadataOuterPuzzle:
-    _match: Any
-    _asset_id: Any
-    _construct: Any
-    _solve: Any
-    _get_inner_puzzle: Any
-    _get_inner_solution: Any
+    _match: Callable[[Program], Optional[PuzzleInfo]]
+    _construct: Callable[[PuzzleInfo, Program], Program]
+    _solve: Callable[[PuzzleInfo, Solver, Program, Program], Program]
+    _get_inner_puzzle: Callable[[PuzzleInfo, Program], Optional[Program]]
+    _get_inner_solution: Callable[[PuzzleInfo, Program], Optional[Program]]
 
     def match(self, puzzle: Program) -> Optional[PuzzleInfo]:
         matched, curried_args = match_metadata_layer_puzzle(puzzle)
@@ -59,16 +58,18 @@ class MetadataOuterPuzzle:
         return bytes32(constructor["updater_hash"])
 
     def construct(self, constructor: PuzzleInfo, inner_puzzle: Program) -> Program:
-        if constructor.also() is not None:
-            inner_puzzle = self._construct(constructor.also(), inner_puzzle)
+        also = constructor.also()
+        if also is not None:
+            inner_puzzle = self._construct(also, inner_puzzle)
         return puzzle_for_metadata_layer(constructor["metadata"], constructor["updater_hash"], inner_puzzle)
 
     def get_inner_puzzle(self, constructor: PuzzleInfo, puzzle_reveal: Program) -> Optional[Program]:
         matched, curried_args = match_metadata_layer_puzzle(puzzle_reveal)
         if matched:
             _, _, _, inner_puzzle = curried_args
-            if constructor.also() is not None:
-                deep_inner_puzzle: Optional[Program] = self._get_inner_puzzle(constructor.also(), inner_puzzle)
+            also = constructor.also()
+            if also is not None:
+                deep_inner_puzzle: Optional[Program] = self._get_inner_puzzle(also, inner_puzzle)
                 return deep_inner_puzzle
             else:
                 return inner_puzzle
@@ -77,8 +78,9 @@ class MetadataOuterPuzzle:
 
     def get_inner_solution(self, constructor: PuzzleInfo, solution: Program) -> Optional[Program]:
         my_inner_solution: Program = solution.first()
-        if constructor.also():
-            deep_inner_solution: Optional[Program] = self._get_inner_solution(constructor.also(), my_inner_solution)
+        also = constructor.also()
+        if also:
+            deep_inner_solution: Optional[Program] = self._get_inner_solution(also, my_inner_solution)
             return deep_inner_solution
         else:
             return my_inner_solution
@@ -86,8 +88,9 @@ class MetadataOuterPuzzle:
     def solve(self, constructor: PuzzleInfo, solver: Solver, inner_puzzle: Program, inner_solution: Program) -> Program:
         coin_bytes: bytes = solver["coin"]
         coin: Coin = Coin(bytes32(coin_bytes[0:32]), bytes32(coin_bytes[32:64]), uint64.from_bytes(coin_bytes[64:72]))
-        if constructor.also() is not None:
-            inner_solution = self._solve(constructor.also(), solver, inner_puzzle, inner_solution)
+        also = constructor.also()
+        if also is not None:
+            inner_solution = self._solve(also, solver, inner_puzzle, inner_solution)
         return solution_for_metadata_layer(
             uint64(coin.amount),
             inner_solution,

chia/wallet/nft_wallet/ownership_outer_puzzle.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, List, Optional, Tuple, Union
+from typing import Callable, List, Optional, Tuple, Union
 
 from clvm_tools.binutils import disassemble
 
@@ -30,12 +30,11 @@ def solution_for_ownership_layer(inner_solution: Program) -> Program:
 
 @dataclass(frozen=True)
 class OwnershipOuterPuzzle:
-    _match: Any
-    _asset_id: Any
-    _construct: Any
-    _solve: Any
-    _get_inner_puzzle: Any
-    _get_inner_solution: Any
+    _match: Callable[[Program], Optional[PuzzleInfo]]
+    _construct: Callable[[PuzzleInfo, Program], Program]
+    _solve: Callable[[PuzzleInfo, Solver, Program, Program], Program]
+    _get_inner_puzzle: Callable[[PuzzleInfo, Program], Optional[Program]]
+    _get_inner_solution: Callable[[PuzzleInfo, Program], Optional[Program]]
 
     def match(self, puzzle: Program) -> Optional[PuzzleInfo]:
         matched, curried_args = match_ownership_layer_puzzle(puzzle)
@@ -61,8 +60,9 @@ class OwnershipOuterPuzzle:
         return None
 
     def construct(self, constructor: PuzzleInfo, inner_puzzle: Program) -> Program:
-        if constructor.also() is not None:
-            inner_puzzle = self._construct(constructor.also(), inner_puzzle)
+        also = constructor.also()
+        if also is not None:
+            inner_puzzle = self._construct(also, inner_puzzle)
         transfer_program_info: Union[PuzzleInfo, Program] = constructor["transfer_program"]
         if isinstance(transfer_program_info, Program):
             transfer_program: Program = transfer_program_info
@@ -74,8 +74,9 @@ class OwnershipOuterPuzzle:
         matched, curried_args = match_ownership_layer_puzzle(puzzle_reveal)
         if matched:
             _, _, _, inner_puzzle = curried_args
-            if constructor.also() is not None:
-                deep_inner_puzzle: Optional[Program] = self._get_inner_puzzle(constructor.also(), inner_puzzle)
+            also = constructor.also()
+            if also is not None:
+                deep_inner_puzzle: Optional[Program] = self._get_inner_puzzle(also, inner_puzzle)
                 return deep_inner_puzzle
             else:
                 return inner_puzzle
@@ -84,13 +85,15 @@ class OwnershipOuterPuzzle:
 
     def get_inner_solution(self, constructor: PuzzleInfo, solution: Program) -> Optional[Program]:
         my_inner_solution: Program = solution.first()
-        if constructor.also():
-            deep_inner_solution: Optional[Program] = self._get_inner_solution(constructor.also(), my_inner_solution)
+        also = constructor.also()
+        if also:
+            deep_inner_solution: Optional[Program] = self._get_inner_solution(also, my_inner_solution)
             return deep_inner_solution
         else:
             return my_inner_solution
 
     def solve(self, constructor: PuzzleInfo, solver: Solver, inner_puzzle: Program, inner_solution: Program) -> Program:
-        if constructor.also() is not None:
-            inner_solution = self._solve(constructor.also(), solver, inner_puzzle, inner_solution)
+        also = constructor.also()
+        if also is not None:
+            inner_solution = self._solve(also, solver, inner_puzzle, inner_solution)
         return solution_for_ownership_layer(inner_solution)

chia/wallet/nft_wallet/singleton_outer_puzzle.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Callable, Optional
 
 from chia.types.blockchain_format.coin import Coin
 from chia.types.blockchain_format.program import Program
@@ -18,12 +18,11 @@ from chia.wallet.puzzles.singleton_top_layer_v1_1 import (
 
 @dataclass(frozen=True)
 class SingletonOuterPuzzle:
-    _match: Any
-    _asset_id: Any
-    _construct: Any
-    _solve: Any
-    _get_inner_puzzle: Any
-    _get_inner_solution: Any
+    _match: Callable[[Program], Optional[PuzzleInfo]]
+    _construct: Callable[[PuzzleInfo, Program], Program]
+    _solve: Callable[[PuzzleInfo, Solver, Program, Program], Program]
+    _get_inner_puzzle: Callable[[PuzzleInfo, Program], Optional[Program]]
+    _get_inner_solution: Callable[[PuzzleInfo, Program], Optional[Program]]
 
     def match(self, puzzle: Program) -> Optional[PuzzleInfo]:
         matched, curried_args = match_singleton_puzzle(puzzle)
@@ -45,8 +44,9 @@ class SingletonOuterPuzzle:
         return bytes32(constructor["launcher_id"])
 
     def construct(self, constructor: PuzzleInfo, inner_puzzle: Program) -> Program:
-        if constructor.also() is not None:
-            inner_puzzle = self._construct(constructor.also(), inner_puzzle)
+        also = constructor.also()
+        if also is not None:
+            inner_puzzle = self._construct(also, inner_puzzle)
         launcher_hash = constructor["launcher_ph"] if "launcher_ph" in constructor else SINGLETON_LAUNCHER_HASH
         return puzzle_for_singleton(constructor["launcher_id"], inner_puzzle, launcher_hash)
 
@@ -54,8 +54,9 @@ class SingletonOuterPuzzle:
         matched, curried_args = match_singleton_puzzle(puzzle_reveal)
         if matched:
             _, inner_puzzle = curried_args
-            if constructor.also() is not None:
-                deep_inner_puzzle: Optional[Program] = self._get_inner_puzzle(constructor.also(), inner_puzzle)
+            also = constructor.also()
+            if also is not None:
+                deep_inner_puzzle: Optional[Program] = self._get_inner_puzzle(also, inner_puzzle)
                 return deep_inner_puzzle
             else:
                 return inner_puzzle
@@ -64,8 +65,9 @@ class SingletonOuterPuzzle:
 
     def get_inner_solution(self, constructor: PuzzleInfo, solution: Program) -> Optional[Program]:
         my_inner_solution: Program = solution.at("rrf")
-        if constructor.also():
-            deep_inner_solution: Optional[Program] = self._get_inner_solution(constructor.also(), my_inner_solution)
+        also = constructor.also()
+        if also:
+            deep_inner_solution: Optional[Program] = self._get_inner_solution(also, my_inner_solution)
             return deep_inner_solution
         else:
             return my_inner_solution
@@ -75,8 +77,9 @@ class SingletonOuterPuzzle:
         coin: Coin = Coin(bytes32(coin_bytes[0:32]), bytes32(coin_bytes[32:64]), uint64.from_bytes(coin_bytes[64:72]))
         parent_spend: CoinSpend = CoinSpend.from_bytes(solver["parent_spend"])
         parent_coin: Coin = parent_spend.coin
-        if constructor.also() is not None:
-            inner_solution = self._solve(constructor.also(), solver, inner_puzzle, inner_solution)
+        also = constructor.also()
+        if also is not None:
+            inner_solution = self._solve(also, solver, inner_puzzle, inner_solution)
         matched, curried_args = match_singleton_puzzle(parent_spend.puzzle_reveal.to_program())
         assert matched
         _, parent_inner_puzzle = curried_args

chia/wallet/nft_wallet/transfer_program_puzzle.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, List, Optional, Tuple
+from typing import Callable, List, Optional, Tuple
 
 from chia.types.blockchain_format.program import Program
 from chia.types.blockchain_format.sized_bytes import bytes32
@@ -39,12 +39,11 @@ def solution_for_transfer_program(
 
 @dataclass(frozen=True)
 class TransferProgramPuzzle:
-    _match: Any
-    _asset_id: Any
-    _construct: Any
-    _solve: Any
-    _get_inner_puzzle: Any
-    _get_inner_solution: Any
+    _match: Callable[[Program], Optional[PuzzleInfo]]
+    _construct: Callable[[PuzzleInfo, Program], Program]
+    _solve: Callable[[PuzzleInfo, Solver, Program, Program], Program]
+    _get_inner_puzzle: Callable[[PuzzleInfo, Program], Optional[Program]]
+    _get_inner_solution: Callable[[PuzzleInfo, Program], Optional[Program]]
 
     def match(self, puzzle: Program) -> Optional[PuzzleInfo]:
         matched, curried_args = match_transfer_program_puzzle(puzzle)

chia/wallet/outer_puzzles.py

@@ -63,11 +63,11 @@ def get_inner_solution(constructor: PuzzleInfo, solution: Program) -> Optional[P
     return driver_lookup[AssetType(constructor.type())].get_inner_solution(constructor, solution)
 
 
-def create_asset_id(constructor: PuzzleInfo) -> bytes32:
-    return driver_lookup[AssetType(constructor.type())].asset_id(constructor)  # type: ignore
+def create_asset_id(constructor: PuzzleInfo) -> Optional[bytes32]:
+    return driver_lookup[AssetType(constructor.type())].asset_id(constructor)
 
 
-function_args = [match_puzzle, create_asset_id, construct_puzzle, solve_puzzle, get_inner_puzzle, get_inner_solution]
+function_args = (match_puzzle, construct_puzzle, solve_puzzle, get_inner_puzzle, get_inner_solution)
 
 driver_lookup: Dict[AssetType, DriverProtocol] = {
     AssetType.CAT: CATOuterPuzzle(*function_args),

chia/wallet/puzzle_drivers.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 
@@ -51,7 +53,7 @@ class PuzzleInfo:
     def type(self) -> str:
         return str(self.info["type"])
 
-    def also(self) -> Optional["PuzzleInfo"]:
+    def also(self) -> Optional[PuzzleInfo]:
         if "also" in self.info:
             return PuzzleInfo(self.info["also"])
         else:

tests/wallet/cat_wallet/test_cat_outer_puzzle.py

@@ -9,14 +9,7 @@ from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.types.coin_spend import CoinSpend
 from chia.util.ints import uint64
 from chia.wallet.cat_wallet.cat_utils import construct_cat_puzzle
-from chia.wallet.outer_puzzles import (
-    construct_puzzle,
-    create_asset_id,
-    get_inner_puzzle,
-    get_inner_solution,
-    match_puzzle,
-    solve_puzzle,
-)
+from chia.wallet.outer_puzzles import construct_puzzle, get_inner_puzzle, get_inner_solution, match_puzzle, solve_puzzle
 from chia.wallet.puzzle_drivers import PuzzleInfo, Solver
 from chia.wallet.puzzles.cat_loader import CAT_MOD
 
@@ -37,7 +30,6 @@ def test_cat_outer_puzzle() -> None:
     assert inside_cat_driver["tail"] == tail
     assert construct_puzzle(cat_driver, ACS) == double_cat_puzzle
     assert get_inner_puzzle(cat_driver, double_cat_puzzle) == ACS
-    assert create_asset_id(cat_driver) == tail
 
     # Set up for solve
     parent_coin = Coin(tail, double_cat_puzzle.get_tree_hash(), uint64(100))

tests/wallet/nft_wallet/test_ownership_outer_puzzle.py

@@ -7,14 +7,7 @@ from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.util.ints import uint16
 from chia.wallet.nft_wallet.ownership_outer_puzzle import puzzle_for_ownership_layer
 from chia.wallet.nft_wallet.transfer_program_puzzle import puzzle_for_transfer_program
-from chia.wallet.outer_puzzles import (
-    construct_puzzle,
-    create_asset_id,
-    get_inner_puzzle,
-    get_inner_solution,
-    match_puzzle,
-    solve_puzzle,
-)
+from chia.wallet.outer_puzzles import construct_puzzle, get_inner_puzzle, get_inner_solution, match_puzzle, solve_puzzle
 from chia.wallet.puzzle_drivers import PuzzleInfo, Solver
 
 
@@ -56,7 +49,6 @@ def test_ownership_outer_puzzle() -> None:
     assert construct_puzzle(ownership_driver_empty, ACS) == ownership_puzzle_empty
     assert construct_puzzle(ownership_driver_default, ACS) == ownership_puzzle_default
     assert get_inner_puzzle(ownership_driver, ownership_puzzle) == ACS
-    assert create_asset_id(ownership_driver) is None
 
     # Set up for solve
     inner_solution = Program.to(