pymake

commit 2dfbf1fea450d17f8a3da8e67f248ff9f7b6548f
parent 81f44ad2bc289f1f77b5854c1bad4a7d20755fbd
Author: gracefu <81774659+gracefuu@users.noreply.github.com>
Date:   Wed, 16 Apr 2025 00:25:17 +0800

Fix types, remove rawrule and add a function to register rules with Rules

Diffstat:
M
make.py
 | 
47
++++++++---------------------------------------

1 file changed, 8 insertions(+), 39 deletions(-)
diff --git a/make.py b/make.py
@@ -55,12 +55,13 @@ import hashlib
 
 from typing import (
     Any,
+    Awaitable,
+    Concatenate,
     Optional,
     Callable,
     Protocol,
     Tuple,
     List,
-    TypedDict,
 )
 
 
@@ -75,40 +76,8 @@ TaskKey = tuple
 ValueHash = bytes
 TaskInputs = List[Tuple[TaskKey, ValueHash]]
 
-
-class RuleFn(Protocol):
-    """Protocol for rule functions that produce task values."""
-
-    async def __call__(
-        self,
-        fetch: Fetch,
-        task_key: TaskKey,
-        store: "Store",
-        *args: Any,
-    ) -> Any: ...
-
-
-class NiceRuleFn(Protocol):
-    """Protocol for simplified rule functions that produce task values."""
-
-    async def __call__(
-        self,
-        fetch: Fetch,
-        *args: Any,
-    ) -> Any: ...
-
-
-class CacheFn(Protocol):
-    """Protocol for cache functions that call rule functions."""
-
-    async def __call__(
-        self,
-        fetch: Fetch,
-        task_key: TaskKey,
-        store: "Store",
-        rule_fn: RuleFn,
-        *args: Any,
-    ) -> Any: ...
+RuleFn = Callable[Concatenate[Fetch, TaskKey, "Store", ...], Awaitable[Any]]
+NiceRuleFn = Callable[Concatenate[Fetch, ...], Awaitable[Any]]
 
 
 def make_hash(o: Any) -> bytes:
@@ -223,16 +192,16 @@ class Rules:
         return rule(*args)
 
     def rule(self, rule_fn: NiceRuleFn) -> Rule:
+        @self.register
         @self.hash_cache
-        @self.rawrule
+        @Rule.new
         @functools.wraps(rule_fn)
         def wrapped(fetch, task_key, store, *args):
             return rule_fn(fetch, *args)
 
         return wrapped
 
-    def rawrule(self, rule_fn: RuleFn) -> Rule:
-        rule = Rule.new(rule_fn)
+    def register(self, rule: Rule) -> Rule:
         self.rules[rule.rule_key] = rule
         return rule
 
@@ -295,7 +264,7 @@ class Rules:
 
 _rules = Rules()
 rule = _rules.rule
-rawrule = _rules.rawrule
+register = _rules.register
 hash_cache = _rules.hash_cache