pymake

commit 7fc5bfcf341c686524755e71d8e835833134a544
parent e6a64f5972b99b6fa6a8bf0f0a879449fbde9bd6
Author: gracefu <81774659+gracefuu@users.noreply.github.com>
Date:   Tue, 15 Apr 2025 13:19:54 +0800

Minor typing tweaks

Diffstat:
M
examples.py
 | 
2
+-
M
make.py
 | 
24
+++++++++++++-----------

2 files changed, 14 insertions(+), 12 deletions(-)
diff --git a/examples.py b/examples.py
@@ -84,7 +84,7 @@ async def run_examples():
         # Note that `build(...)` will wait for all detached jobs to complete before returning.
         # You may choose to use the lower level `build.fetch(...)` function instead, which does not wait for detached jobs.
         # You must then ensure `build.wait()` is called later to wait for detached jobs to complete.
-        await build.fetch(_eg_rec(2345))
+        detach(build.fetch(_eg_rec(2345)))
         await build.fetch(_eg_rec(3456))
         await build.wait()
 
diff --git a/make.py b/make.py
@@ -35,8 +35,10 @@ from typing import Awaitable, Callable, Any, Concatenate, Optional
 
 
 Fetch = Callable[["Task"], Awaitable[Any]]
-TaskKey = str
-RuleKey = str
+TaskKey = tuple
+RuleKey = bytes
+
+RuleFn = Callable[Concatenate[Fetch, TaskKey, "Store", ...], Awaitable[Any]]
 
 
 def _make_hash(o: Any) -> bytes:
@@ -50,17 +52,17 @@ def _make_hash(o: Any) -> bytes:
     return h.digest()
 
 
-def _fn_to_key(fn) -> str:
+def _rule_fn_to_key(fn) -> RuleKey:
     name = fn.__name__
     source = inspect.getsource(fn)
     h = hashlib.sha256(source.encode("utf-8")).hexdigest()[:16]
-    key = f"{name}-{len(source)}-{h}"
+    key = f"{name}-{len(source)}-{h}".encode("utf-8")
     return key
 
 
 class Task:
     task_key: TaskKey
-    rule_fn: Callable[Concatenate[Fetch, TaskKey, "Store", ...], Awaitable[Any]]
+    rule_fn: RuleFn
     args: tuple
     hash: int
 
@@ -69,7 +71,7 @@ class Task:
         return Task(
             (
                 rule.rule_key,
-                *(arg.task_key if hasattr(arg, "task_key") else arg for arg in args),
+                *(arg.task_key if isinstance(arg, Task) else arg for arg in args),
             ),
             rule.rule_fn,
             *args,
@@ -96,14 +98,14 @@ class Task:
 
 class Rule:
     rule_key: RuleKey
-    rule_fn: Callable[Concatenate[Fetch, TaskKey, "Store", ...], Awaitable[Any]]
+    rule_fn: RuleFn
     hash: int
 
     @staticmethod
-    def new(rule_fn):
-        return Rule(_fn_to_key(rule_fn), rule_fn)
+    def new(rule_fn: RuleFn):
+        return Rule(_rule_fn_to_key(rule_fn), rule_fn)
 
-    def __init__(self, rule_key, rule_fn):
+    def __init__(self, rule_key: RuleKey, rule_fn: RuleFn):
         self.rule_key = rule_key
         self.rule_fn = rule_fn
         self.hash = hash(self.rule_key)
@@ -153,7 +155,7 @@ class Rules:
     def cache(self):
         def decorator(rule: Rule):
             @functools.wraps(rule.rule_fn)
-            async def new_rule_fn(fetch: Fetch, task_key: str, store: "Store", *args):
+            async def new_rule_fn(fetch: Fetch, task_key: TaskKey, store: Store, *args):
                 past_runs = store.key_info[task_key]
                 output_value = store.key_value[task_key]
                 possible_values = []