pymake

__init__.py (7689B)

---

 import asyncio
 from asyncio.futures import Future
 from contextlib import asynccontextmanager
 from contextvars import ContextVar
 from types import CoroutineType
 from typing import (
     Any,
     Awaitable,
     Callable,
     ClassVar,
     Generator,
     Generic,
     Protocol,
     TypeAlias,
     TypeVar,
     runtime_checkable,
 )
 import dataclasses
 
 # Alias confusing names that clash with the build system
 AIOTask = asyncio.Task
 AIOTaskGroup = asyncio.TaskGroup
 AIOEvent = asyncio.Event
 
 
 CovT = TypeVar("CovT", covariant=True)
 ConT = TypeVar("ConT", contravariant=True)
 
 DbEntry: TypeAlias = object
 
 Factory: TypeAlias = Callable[[], CovT]
 
 
 @runtime_checkable
 class HashLike(Protocol):
     def __hash__(self) -> int: ...
     def __eq__(self, other, /) -> bool: ...
 
 
 @dataclasses.dataclass(frozen=True)
 class TaskKey:
     rule_key: str
     arg_keys: tuple["TaskKey | HashLike"]
 
 
 @dataclasses.dataclass(slots=True)
 class Box(Generic[CovT]):
     x: CovT = dataclasses.field()
 
 
 build_tg_: ContextVar[AIOTaskGroup] = ContextVar("Build asyncio.TaskGroup")
 build_seq_: ContextVar[Box[int]] = ContextVar("Build sequence number")
 
 
 def build_tg():
     return build_tg_.get()
 
 
 def build_seq():
     seq = build_seq_.get()
     seq.x += 1
     return seq.x
 
 
 @asynccontextmanager
 async def BuildContext():
     async with AIOTaskGroup() as tg:
         tok = build_tg_.set(tg)
         build_seq_.set(Box(0))
         try:
             yield
         finally:
             build_tg_.reset(tok)
 
 
 @dataclasses.dataclass
 class TaskContext:
     task: "Task" = dataclasses.field()
 
 
 async def noop():
     pass
 
 
 @dataclasses.dataclass(slots=True)
 class LazyTask(Generic[CovT]):
     """
     Lazy AIOTask
 
     Lazy: The coroutine isn't run or spawned into a task until .start()ed or awaited.
     Use await directly if you intend to await -- this runs the coroutine in the current task.
     """
 
     coro: CoroutineType[Any, None, CovT]
     # Not started | Starting | Started
     task: None | Future[CovT] = None
 
     def start(self) -> Awaitable[CovT]:
         """Ensures task is started, and returns a coroutine for its completion."""
         if self.task is None:
             self.task = build_tg().create_task(self.coro)
         return self.task
 
     def __await__(self) -> Generator[Any, None, CovT]:
         if self.task is None:
             # Run the coroutine directly to avoid creating an additional task.
             self.task = asyncio.get_event_loop().create_future()
             try:
                 res = yield from self.coro.__await__()
                 self.task.set_result(res)
                 return res
             except BaseException as exc:
                 self.task.set_exception(exc)
                 raise
         else:
             res = yield from self.task.__await__()
             return res
 
 
 @dataclasses.dataclass(slots=True, frozen=True)
 class LazyFuture(Generic[CovT]):
     """
     Future that lazily starts its own computation.
 
     The handle is run the first time .start() is called or the future is awaited.
 
     For convenience, LazyFutures are also context managers that intercept exceptions.
     """
 
     handle: LazyTask | None = dataclasses.field(default=None)
     fut: Future[CovT] = dataclasses.field(
         default_factory=lambda: asyncio.get_event_loop().create_future()
     )
 
     none: ClassVar["LazyFuture[None]"]
 
     def __enter__(self):
         return self.fut.set_result
 
     def __exit__(self, exc_type, exc, tb):
         if not self.fut.done():
             self.fut.set_exception(exc)
 
     def start(self):
         """Ensures handle is started if any, then returns the future."""
         if self.handle:
             self.handle.start()
         return self.fut
 
     def __await__(self) -> Generator[Any, None, CovT]:
         """Ensures handle is started if any, then awaits the future."""
         if self.handle is not None:
             yield from self.handle.start().__await__()
         res = yield from self.fut.__await__()
         return res
 
 
 LazyFuture.none = LazyFuture()
 LazyFuture.none.fut.set_result(None)
 
 
 CovT = TypeVar("CovT", covariant=True)
 
 
 class Unset: ...
 
 
 unset = Unset()
 
 
 @dataclasses.dataclass(frozen=True)
 class Task(Generic[CovT]):
     """
     A task is a computation of the task's value and hash.
 
     Values are returned when other tasks fetch this task.
 
     Other tasks can also obtain this task's hash, but hashes should be viewed
     as a way to certify that "a value" is the same as some "other value",
     typically to assert that the value from a previous run is good.
     """
 
     key: TaskKey = dataclasses.field()
     value: LazyFuture[CovT] = dataclasses.field(compare=False)
     hash: LazyFuture[HashLike] = dataclasses.field(compare=False)
 
 
 @dataclasses.dataclass
 class TaskRun(Generic[CovT]):
     """
     A task run is a recording of how the task ran last time.
     """
 
     key: TaskKey = dataclasses.field()
     deps: dict[TaskKey, "TaskRun"] = dataclasses.field(default_factory=dict)
     value: CovT | Unset = dataclasses.field(default=unset)
     hash: HashLike | Unset = dataclasses.field(default=unset)
 
 
 """
 Traces record information of a run so work can be saved in the future.
 
 There are 3 variants, dep traces, hash traces and value traces.
 
 - A dep trace records a task's deps.
 - A hash trace records a task's deps, dep hashes, and *out hash*.
 - A value trace records a task's deps, dep hashes, *out hash, and out value*.
 
 Trace builders are classes to aid in the production of a trace.
 Once a trace builder can be finalized to obtain a true trace.
 
 As such, there are 4 classes:
 
 - DepTrace
 - HashTrace
 - ValueTrace
 - DepTraceBuilder
 - HashTraceBuilder
 - ValueTraceBuilder
 """
 
 
 @dataclasses.dataclass
 class DepTrace:
     key: TaskKey = dataclasses.field()
     # We still record HashLike so HashTrace and ValueTrace can subclass DepTrace.
     # The hash is simply ignored if we only want to know the dependencies of key.
     in_hashes: dict[TaskKey, HashLike] = dataclasses.field()
 
 
 @dataclasses.dataclass
 class HashTrace(DepTrace):
     out_hash: HashLike = dataclasses.field()
 
 
 @dataclasses.dataclass
 class ValueTrace(HashTrace, Generic[CovT]):
     out_value: CovT = dataclasses.field()
 
 
 @dataclasses.dataclass
 class DepTraceBuilder:
     key: TaskKey = dataclasses.field()
     in_hashes: dict[TaskKey, LazyFuture[HashLike]] = dataclasses.field(
         default_factory=dict
     )
 
     async def finalize(self) -> DepTrace:
         return DepTrace(
             key=self.key,
             in_hashes={k: await v for k, v in self.in_hashes.items()},
         )
 
 
 @dataclasses.dataclass
 class HashTraceBuilder(DepTraceBuilder):
     out_hash: LazyFuture[HashLike] | None = dataclasses.field(default=None)
 
     async def finalize(self) -> HashTrace:
         if self.out_hash is None:
             raise RuntimeError("Cannot finalize HashTraceBuilder without out_hash")
         return HashTrace(
             key=self.key,
             in_hashes={k: await v for k, v in self.in_hashes.items()},
             out_hash=await self.out_hash,
         )
 
 
 @dataclasses.dataclass
 class ValueTraceBuilder(HashTraceBuilder, Generic[CovT]):
     out_value: LazyFuture[CovT] | None = dataclasses.field(default=None)
 
     async def finalize(self) -> ValueTrace[CovT]:
         if self.out_hash is None:
             raise RuntimeError("Cannot finalize ValueTraceBuilder without out_hash")
         if self.out_value is None:
             raise RuntimeError("Cannot finalize ValueTraceBuilder without out_value")
         return ValueTrace(
             key=self.key,
             in_hashes={k: await v for k, v in self.in_hashes.items()},
             out_hash=await self.out_hash,
             out_value=await self.out_value,
         )