address Richard's review comments on the recent Program.run() changes

Chia-Network · Jul 31, 2024 · 8a2df86 · 8a2df86
1 parent e5b1e85
commit 8a2df86
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 13 deletions.
diff --git a/chia/_tests/clvm/test_program.py b/chia/_tests/clvm/test_program.py
@@ -119,10 +119,28 @@ def test_run() -> None:
     ret = div.run([10, -5])
     assert ret.atom == bytes([0xFE])
 
+    # run()
+    with pytest.raises(ValueError, match="div operator with negative operands is deprecated"):
+        cost, ret = div.run_with_cost(100000, [10, -5], 0)
+
+    cost, ret = div.run_with_cost(100000, [10, -5], ENABLE_FIXED_DIV)
+    assert cost == 1107
+    print(ret)
+    assert ret.atom == bytes([0xFE])
+
+    # run_with_flags()
     with pytest.raises(ValueError, match="div operator with negative operands is deprecated"):
         cost, ret = div.run_with_flags(100000, 0, [10, -5])
 
     cost, ret = div.run_with_flags(100000, ENABLE_FIXED_DIV, [10, -5])
     assert cost == 1107
     print(ret)
     assert ret.atom == bytes([0xFE])
+
+    # run_with_cost()
+    with pytest.raises(ValueError, match="div operator with negative operands is deprecated"):
+        cost, ret = div.run([10, -5], 100000, 0)
+
+    ret = div.run([10, -5], 100000, ENABLE_FIXED_DIV)
+    print(ret)
+    assert ret.atom == bytes([0xFE])
diff --git a/chia/types/blockchain_format/program.py b/chia/types/blockchain_format/program.py
@@ -29,6 +29,15 @@
 
 INFINITE_COST = 11000000000
 
+DEFAULT_FLAGS = (
+    ENABLE_SOFTFORK_CONDITION
+    | ENABLE_BLS_OPS_OUTSIDE_GUARD
+    | ENABLE_FIXED_DIV
+    | AGG_SIG_ARGS
+    | ENABLE_MESSAGE_CONDITIONS
+    | DISALLOW_INFINITY_G1
+    | MEMPOOL_MODE
+)
 
 T_CLVMStorage = TypeVar("T_CLVMStorage", bound=CLVMStorage)
 T_Program = TypeVar("T_Program", bound="Program")
@@ -139,22 +148,13 @@ def _run(self, max_cost: int, flags: int, args: Any) -> Tuple[int, Program]:
         cost, r = run_chia_program(self.as_bin(), prog_args.as_bin(), max_cost, flags)
         return cost, Program.to(r)
 
-    def run_with_cost(self, max_cost: int, args: Any) -> Tuple[int, Program]:
+    def run_with_cost(self, max_cost: int, args: Any, flags=DEFAULT_FLAGS) -> Tuple[int, Program]:
         # when running puzzles in the wallet, default to enabling all soft-forks
         # as well as enabling mempool-mode (i.e. strict mode)
-        default_flags = (
-            ENABLE_SOFTFORK_CONDITION
-            | ENABLE_BLS_OPS_OUTSIDE_GUARD
-            | ENABLE_FIXED_DIV
-            | AGG_SIG_ARGS
-            | ENABLE_MESSAGE_CONDITIONS
-            | DISALLOW_INFINITY_G1
-            | MEMPOOL_MODE
-        )
-        return self._run(max_cost, default_flags, args)
+        return self._run(max_cost, flags, args)
 
-    def run(self, args: Any) -> Program:
-        cost, r = self.run_with_cost(INFINITE_COST, args)
+    def run(self, args: Any, max_cost=INFINITE_COST, flags=DEFAULT_FLAGS) -> Program:
+        cost, r = self._run(max_cost, flags, args)
         return r
 
     def run_with_flags(self, max_cost: int, flags: int, args: Any) -> Tuple[int, Program]: