commit 36d24d478d058203fc16309340adbef75cc411ed
parent c376a67029ba96db031296a5d80a3d94a6dd2c32
Author: gracefu <81774659+gracefuu@users.noreply.github.com>
Date: Mon, 21 Apr 2025 01:07:37 +0800
Replace file pickling with open wrapper
Diffstat:
7 files changed, 59 insertions(+), 51 deletions(-)
diff --git a/make3/__init__.py b/make3/__init__.py
@@ -1,3 +1,4 @@
+from .io import *
from .pickler import *
from .rebuild import *
from .helpers import *
diff --git a/make3/helpers.py b/make3/helpers.py
@@ -1,4 +1,4 @@
-from typing import IO, Any
+from .io import open
from .rebuild import cache_conditionally, rerun_if_changed
import hashlib
import os
@@ -9,7 +9,7 @@ def file_modtime(f: int | str | bytes | os.PathLike[str] | os.PathLike[bytes]):
@cache_conditionally(lambda f, *args: (f.name, *args))
-def _file_hash(f: IO[Any], skip_if_modtime_matches=True):
+def _file_hash(f: open, skip_if_modtime_matches=True):
if skip_if_modtime_matches:
rerun_if_changed(lambda: file_modtime(f.fileno()))
else:
@@ -22,7 +22,7 @@ def _file_hash(f: IO[Any], skip_if_modtime_matches=True):
return d
-def file_hash(f: IO[Any] | bytes | str, skip_if_modtime_matches=True):
+def file_hash(f: open | 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)
diff --git a/make3/io.py b/make3/io.py
@@ -0,0 +1,34 @@
+_open = open
+
+
+class open:
+ def __init__(self, *args, **kwargs):
+ self.args = args
+ self.kwargs = kwargs
+ self.entered = False
+ self.file = _open(*args, **kwargs)
+
+ @staticmethod
+ def __open_seek(pos, entered, args, kwargs):
+ f = open(*args, **kwargs)
+ if entered:
+ f.__enter__()
+ f.seek(pos)
+ return f
+
+ def __getattr__(self, attr):
+ return getattr(self.file, attr)
+
+ def __iter__(self):
+ return iter(self.file)
+
+ def __enter__(self):
+ self.file.__enter__()
+ self.entered = True
+ return self
+
+ def __exit__(self, ty, exc, tb):
+ self.file.__exit__(ty, exc, tb)
+
+ def __reduce__(self):
+ return (self.__open_seek, (self.tell(), self.entered, self.args, self.kwargs))
diff --git a/make3/pickler.py b/make3/pickler.py
@@ -1,7 +1,6 @@
-# from io import BufferedReader, IOBase
-from _io import _IOBase, BufferedReader
-from io import BufferedIOBase, BufferedRandom, BufferedWriter
+from io import BytesIO
from types import CellType, CodeType, FunctionType
+from typing import Any
import copyreg
import importlib
import inspect
@@ -75,9 +74,6 @@ class BasePickler(pickle.Pickler):
return NotImplemented
-# BufferedReader
-
-
def FunctionPicklerAddon(pickler: type[BasePickler]):
class Wrapper(pickler):
dispatch_table = pickler.dispatch_table.copy()
@@ -97,39 +93,12 @@ def FunctionPicklerAddon(pickler: type[BasePickler]):
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)
+
+
+def pickle_with(pickler_cls: type, obj: Any) -> bytes:
+ i = BytesIO()
+ pickler_cls(i).dump(obj)
+ i.seek(0)
+ return i.read()
diff --git a/make3/rebuild.py b/make3/rebuild.py
@@ -1,9 +1,8 @@
+from .pickler import MakePickler, pickle_with
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]]] = (
@@ -11,6 +10,13 @@ rerun_changes_var: contextvars.ContextVar[list[tuple[Any, bytes]]] = (
)
+def get_or_set[T](dict_: dict[str, T], key: str, val: T) -> T:
+ if key in dict_:
+ return dict_[key]
+ dict_[key] = val
+ return val
+
+
def with_rerun_context(rerun_changes, f, /, *args, **kwargs):
rerun_changes_var.set(rerun_changes)
return f(*args, **kwargs)
@@ -21,10 +27,8 @@ 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()))
+ later_pkl = pickle_with(pickler_cls, now if later is None else later)
+ rerun_changes_var.get().append((now() if later is None else now, later_pkl))
def rerun_if(f):
diff --git a/tar-sketch/a.txt b/tar-sketch/a.txt
@@ -1 +1 @@
-Sun Apr 20 08:12:24 PM +08 2025
+Mon Apr 21 01:34:11 AM +08 2025
diff --git a/tar-sketch/tar2.py b/tar-sketch/tar2.py
@@ -1,7 +1,7 @@
import sys
sys.path.append("..")
-from make3 import cache_conditionally, rerun_if_changed, Rerunner, file_hash
+from make3 import cache_conditionally, file_hash, open, rerun_if_changed, Rerunner
import subprocess