commit c376a67029ba96db031296a5d80a3d94a6dd2c32
parent 9379a82f3e413a41335a642081018fb08a9348f7
Author: gracefu <81774659+gracefuu@users.noreply.github.com>
Date: Sun, 20 Apr 2025 20:07:56 +0800
Change rerun syntax to use lambdas instead of strings, enable support for pickling local functions to make it work, added pickling of files because why not, and split package into multiple files
Diffstat:
7 files changed, 278 insertions(+), 165 deletions(-)
diff --git a/make3.py b/make3.py
@@ -1,136 +0,0 @@
-import contextvars
-import functools
-import inspect
-import pickle
-from typing import Any
-
-rerun_db_var: contextvars.ContextVar[dict] = contextvars.ContextVar("rerun_db")
-rerun_changes_var: contextvars.ContextVar[list[tuple[str, Any]]] = (
- contextvars.ContextVar("rerun_changes")
-)
-rerun_globals_var: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar(
- "rerun_globals"
-)
-rerun_locals_var: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar(
- "rerun_locals"
-)
-
-
-def with_rerun_context(
- rerun_changes, rerun_globals, rerun_locals, f, /, *args, **kwargs
-):
- rerun_changes_var.set(rerun_changes)
- rerun_globals_var.set(rerun_globals)
- rerun_locals_var.set(rerun_locals)
- return f(*args, **kwargs)
-
-
-class UseEval: ...
-
-
-def rerun_if_changed(f_str, current_value: Any = UseEval):
- rerun_changes_var.get().append(
- (
- f_str,
- (
- eval(
- f_str,
- globals=rerun_globals_var.get(),
- locals=rerun_locals_var.get(),
- )
- if current_value == UseEval
- else current_value
- ),
- )
- )
-
-
-def rerun_if(f_str):
- return rerun_if_changed(f"bool({f_str})", False)
-
-
-def cache_conditionally(
- keys_fn=lambda *args, **kwargs: (args, tuple(sorted(kwargs.items()))),
- store_fn=lambda result, /, *_, **__: result,
- load_fn=lambda cached_result, /, *_, **__: cached_result,
-):
- def decorator(fn):
- signature = inspect.signature(fn)
-
- @functools.wraps(fn)
- def wrapped(*args, **kwargs):
- db = rerun_db_var.get()
- keys = keys_fn(*args, **kwargs)
- bound_args = signature.bind(*args, **kwargs)
- bound_args.apply_defaults()
- rerun_locals = bound_args.arguments
- if ("track", "result", fn.__qualname__, keys) in db:
- if ("track", "rerun_changes", fn.__qualname__, keys) not in db:
- old_rerun_changes = []
- db[("track", "rerun_changes", fn.__qualname__, keys)] = (
- old_rerun_changes
- )
- else:
- old_rerun_changes = db[
- ("track", "rerun_changes", fn.__qualname__, keys)
- ]
- for expr, old_val in old_rerun_changes:
- try:
- res = eval(expr, globals=fn.__globals__, locals=rerun_locals)
- if res != old_val:
- break
- except BaseException:
- break
- else:
- return load_fn(
- db[("track", "result", fn.__qualname__, keys)], *args, **kwargs
- )
-
- context = contextvars.copy_context()
- rerun_changes = []
- result = context.run(
- with_rerun_context,
- rerun_changes,
- fn.__globals__,
- rerun_locals,
- fn,
- *args,
- **kwargs,
- )
- db[("track", "rerun_changes", fn.__qualname__, keys)] = rerun_changes
- db[("track", "result", fn.__qualname__, keys)] = store_fn(
- result, *args, **kwargs
- )
- return result
-
- return wrapped
-
- return decorator
-
-
-class Rerunner:
- def __init__(self, db_filename=b".makedb", db_file=None):
- if db_file:
- self.db_file = db_file
- else:
- self.db_file = open(db_filename, "a+b")
- self.db_file.seek(0)
-
- def __enter__(self):
- self.db_file.__enter__()
- try:
- self.db = pickle.load(self.db_file)
- except pickle.PickleError:
- self.db = dict()
- except EOFError:
- self.db = dict()
- self.var_tok = rerun_db_var.set(self.db)
- return self
-
- def __exit__(self, ty, exc, tb):
- rerun_db_var.reset(self.var_tok)
- if exc is None:
- self.db_file.seek(0)
- self.db_file.truncate(0)
- pickle.dump(self.db, self.db_file)
- self.db_file.__exit__(ty, exc, tb)
diff --git a/make3/__init__.py b/make3/__init__.py
@@ -0,0 +1,3 @@
+from .pickler import *
+from .rebuild import *
+from .helpers import *
diff --git a/make3/helpers.py b/make3/helpers.py
@@ -0,0 +1,29 @@
+from typing import IO, Any
+from .rebuild import cache_conditionally, rerun_if_changed
+import hashlib
+import os
+
+
+def file_modtime(f: int | str | bytes | os.PathLike[str] | os.PathLike[bytes]):
+ return os.stat(f).st_mtime_ns
+
+
+@cache_conditionally(lambda f, *args: (f.name, *args))
+def _file_hash(f: IO[Any], skip_if_modtime_matches=True):
+ if skip_if_modtime_matches:
+ rerun_if_changed(lambda: file_modtime(f.fileno()))
+ else:
+ rerun_if_changed(False, lambda: True) # always rerun
+ h = hashlib.sha256()
+ for chunk in f:
+ h.update(chunk)
+ d = h.hexdigest()
+ # print("hash", f.name, d)
+ return d
+
+
+def file_hash(f: IO[Any] | bytes | str, skip_if_modtime_matches=True):
+ if isinstance(f, bytes) or isinstance(f, str):
+ with open(f, "rb") as _f:
+ return _file_hash(_f, skip_if_modtime_matches)
+ return _file_hash(f, skip_if_modtime_matches)
diff --git a/make3/pickler.py b/make3/pickler.py
@@ -0,0 +1,135 @@
+# from io import BufferedReader, IOBase
+from _io import _IOBase, BufferedReader
+from io import BufferedIOBase, BufferedRandom, BufferedWriter
+from types import CellType, CodeType, FunctionType
+import copyreg
+import importlib
+import inspect
+import pickle
+
+
+def pickle_code_type(code: CodeType):
+ return (
+ unpickle_code_type,
+ (
+ code.co_argcount,
+ code.co_posonlyargcount,
+ code.co_kwonlyargcount,
+ code.co_nlocals,
+ code.co_stacksize,
+ code.co_flags,
+ code.co_code,
+ code.co_consts,
+ code.co_names,
+ code.co_varnames,
+ code.co_filename,
+ code.co_name,
+ code.co_qualname,
+ code.co_firstlineno,
+ code.co_linetable,
+ code.co_exceptiontable,
+ code.co_freevars,
+ code.co_cellvars,
+ ),
+ )
+
+
+def unpickle_code_type(*args):
+ return CodeType(*args)
+
+
+def pickle_cell_type(cell: CellType):
+ return (unpickle_cell_type, (cell.cell_contents,))
+
+
+def unpickle_cell_type(*args):
+ return CellType(*args)
+
+
+def pickle_function_type(f: FunctionType):
+ mod = inspect.getmodule(f)
+ return (
+ unpickle_function_type,
+ (
+ f.__code__,
+ mod.__name__ if mod is not None else None,
+ (
+ tuple(CellType(cell.cell_contents) for cell in f.__closure__)
+ if f.__closure__
+ else None
+ ),
+ ),
+ )
+
+
+def unpickle_function_type(code, mod_name, closure):
+ return FunctionType(
+ code, globals=importlib.import_module(mod_name).__dict__, closure=closure
+ )
+
+
+class BasePickler(pickle.Pickler):
+ dispatch_table = copyreg.dispatch_table.copy()
+
+ def reducer_override(self, _): # type: ignore
+ return NotImplemented
+
+
+# BufferedReader
+
+
+def FunctionPicklerAddon(pickler: type[BasePickler]):
+ class Wrapper(pickler):
+ dispatch_table = pickler.dispatch_table.copy()
+ dispatch_table[CodeType] = pickle_code_type
+ dispatch_table[CellType] = pickle_cell_type
+
+ def reducer_override(self, obj): # type: ignore
+ if type(obj) is not FunctionType:
+ return super().reducer_override(obj)
+ obj_mod = inspect.getmodule(obj)
+ if obj_mod is None:
+ return super().reducer_override(obj)
+ if obj.__name__ in dir(obj_mod):
+ return super().reducer_override(obj)
+ return pickle_function_type(obj)
+
+ return Wrapper
+
+
+def pickle_buffered(obj: BufferedRandom | BufferedReader | BufferedWriter):
+ d = obj.__dict__
+ # print("aaaaaa", obj.name, obj.mode, obj.tell())
+ return (
+ unpickle_buffered,
+ (obj.name, obj.mode, obj.tell()),
+ )
+
+
+def unpickle_buffered(name, mode, tell):
+ f = open(name, mode)
+ f.seek(tell)
+ return f
+ # print(args)
+ # exit()
+ # return iobaseType(*args)
+
+
+def IOPicklerAddon(pickler: type[BasePickler]):
+ class Wrapper(pickler):
+ def reducer_override(self, obj): # type: ignore
+ # print(type(obj))
+ if (
+ isinstance(obj, BufferedRandom)
+ or isinstance(obj, BufferedReader)
+ or isinstance(obj, BufferedWriter)
+ ):
+ return pickle_buffered(obj)
+ return super().reducer_override(obj)
+
+ return Wrapper
+
+
+MakePickler = BasePickler
+MakePickler = FunctionPicklerAddon(MakePickler)
+MakePickler = IOPicklerAddon(MakePickler)
diff --git a/make3/rebuild.py b/make3/rebuild.py
@@ -0,0 +1,106 @@
+from typing import Any, Callable, overload
+import contextvars
+import functools
+import io
+import pickle
+from .pickler import MakePickler
+
+rerun_db_var: contextvars.ContextVar[dict] = contextvars.ContextVar("rerun_db")
+rerun_changes_var: contextvars.ContextVar[list[tuple[Any, bytes]]] = (
+ contextvars.ContextVar("rerun_changes")
+)
+
+
+def with_rerun_context(rerun_changes, f, /, *args, **kwargs):
+ rerun_changes_var.set(rerun_changes)
+ return f(*args, **kwargs)
+
+
+@overload
+def rerun_if_changed(now: Callable, *, pickler_cls: type = MakePickler): ...
+@overload
+def rerun_if_changed(now: Any, later: Callable, *, pickler_cls: type = MakePickler): ...
+def rerun_if_changed(now, later=None, *, pickler_cls=MakePickler):
+ i = io.BytesIO()
+ pickler_cls(i).dump(now if later is None else later)
+ i.seek(0)
+ rerun_changes_var.get().append((now() if later is None else now, i.read()))
+
+
+def rerun_if(f):
+ return rerun_if_changed(False, lambda: bool(f()))
+
+
+def cache_conditionally(
+ keys_fn=lambda *args, **kwargs: (args, tuple(sorted(kwargs.items()))),
+ store_fn=lambda result, /, *_, **__: result,
+ load_fn=lambda cached_result, /, *_, **__: cached_result,
+):
+ def decorator(fn):
+ @functools.wraps(fn)
+ def wrapped(*args, **kwargs):
+ db = rerun_db_var.get()
+ keys = keys_fn(*args, **kwargs)
+ db_key = ("track", fn.__qualname__, keys)
+ if db_key + ("result",) in db:
+ if db_key + ("rerun_changes",) not in db:
+ old_rerun_changes = []
+ db[db_key + ("rerun_changes",)] = old_rerun_changes
+ else:
+ old_rerun_changes = db[db_key + ("rerun_changes",)]
+ for old_val, f_pkl in old_rerun_changes:
+ try:
+ f_unpkled = pickle.loads(f_pkl)
+ val = f_unpkled()
+ if old_val != val:
+ break
+ except BaseException:
+ break
+ else:
+ return load_fn(db[db_key + ("result",)], *args, **kwargs)
+
+ context = contextvars.copy_context()
+ rerun_changes = []
+ result = context.run(
+ with_rerun_context,
+ rerun_changes,
+ fn,
+ *args,
+ **kwargs,
+ )
+ db[db_key + ("rerun_changes",)] = rerun_changes
+ db[db_key + ("result",)] = store_fn(result, *args, **kwargs)
+ return result
+
+ return wrapped
+
+ return decorator
+
+
+class Rerunner:
+ def __init__(self, db_filename=b".makedb", db_file=None):
+ if db_file:
+ self.db_file = db_file
+ else:
+ self.db_file = open(db_filename, "a+b")
+ self.db_file.seek(0)
+
+ def __enter__(self):
+ self.db_file.__enter__()
+ try:
+ self.db = pickle.load(self.db_file)
+ except pickle.PickleError:
+ self.db = dict()
+ except EOFError:
+ self.db = dict()
+ self.var_tok = rerun_db_var.set(self.db)
+ return self
+
+ def __exit__(self, ty, exc, tb):
+ rerun_db_var.reset(self.var_tok)
+ if exc is None:
+ self.db_file.seek(0)
+ self.db_file.truncate(0)
+ pickle.dump(self.db, self.db_file)
+ # MakePickler(self.db_file).dump(self.db)
+ self.db_file.__exit__(ty, exc, tb)
diff --git a/tar-sketch/a.txt b/tar-sketch/a.txt
@@ -1 +1 @@
-Sun Apr 20 08:02:47 AM +08 2025
+Sun Apr 20 08:12:24 PM +08 2025
diff --git a/tar-sketch/tar2.py b/tar-sketch/tar2.py
@@ -1,45 +1,21 @@
import sys
sys.path.append("..")
-from make3 import cache_conditionally, rerun_if_changed, Rerunner
-
-from io import BufferedReader
-import hashlib
-import os
+from make3 import cache_conditionally, rerun_if_changed, Rerunner, file_hash
import subprocess
-file_modtime = lambda f: os.stat(f.fileno()).st_mtime_ns
-
-
-@cache_conditionally(lambda f: f.name)
-def _file_hash(f: BufferedReader):
- rerun_if_changed("file_modtime(f)")
- h = hashlib.sha256()
- for chunk in f:
- h.update(chunk)
- d = h.hexdigest()
- print("hash", f.name, d)
- return d
-
-
-def file_hash(f: BufferedReader | bytes | str):
- if isinstance(f, bytes) or isinstance(f, str):
- with open(f, "rb") as f:
- return _file_hash(f)
- return _file_hash(f)
-
@cache_conditionally()
def tar(manifest=b"manifest", output=b"archive.tar.gz"):
with open(manifest, "rb") as manifest_f:
manifest_lines = manifest_f.read().splitlines()
- rerun_if_changed(f"file_hash(manifest)")
+ rerun_if_changed(lambda: file_hash(manifest_f))
for fname in manifest_lines:
- rerun_if_changed(f"file_hash({repr(fname)})")
+ rerun_if_changed(lambda: file_hash(fname))
print("tar", "cvzf", output, *manifest_lines)
subprocess.run([b"tar", b"cvzf", output, *manifest_lines])
- rerun_if_changed(f"file_hash({repr(output)})")
+ rerun_if_changed(lambda: file_hash(output))
with Rerunner():