python to_be_file (#366)

Author: nedtwigg

python/example-pytest-selfie/tests/to_be_file_test.py

@@ -0,0 +1,24 @@
+from selfie_lib.WriteTracker import (
+    ToBeFileWriteTracker,
+    ToBeFileLazyBytes,
+    recordCall,
+    TypedPath,
+    SnapshotFileLayout,
+)
+from pytest_selfie import FSImplementation
+from pathlib import Path
+import os
+
+
+def test_to_be_file():
+    layout = SnapshotFileLayout(FSImplementation())
+    tracker = ToBeFileWriteTracker()
+
+    # Record the current call stack.
+    call_stack = recordCall(callerFileOnly=True)
+
+    # Create a TypedPath object for the file path
+    file_path = TypedPath(os.path.abspath(Path("test.jpg")))
+
+    # Write byte data to disk using the tracker.
+    tracker.writeToDisk(file_path, b"some byte data", call_stack, layout)

python/pytest-selfie/pytest_selfie/__init__.py

@@ -1 +1,2 @@
 from .SelfieSettingsAPI import SelfieSettingsAPI as SelfieSettingsAPI
+from .plugin import FSImplementation as FSImplementation

python/pytest-selfie/pytest_selfie/plugin.py

@@ -4,6 +4,7 @@
 from typing import ByteString, DefaultDict, List, Optional, Iterator
 
 from selfie_lib.Atomic import AtomicReference
+from selfie_lib.WriteTracker import ToBeFileWriteTracker
 from .SelfieSettingsAPI import SelfieSettingsAPI
 from selfie_lib import (
     _clearSelfieSystem,
@@ -149,7 +150,7 @@ def __init__(self, settings: SelfieSettingsAPI):
         self._layout = PytestSnapshotFileLayout(self.__fs, settings)
         self.__comment_tracker = CommentTracker()
         self.__inline_write_tracker = InlineWriteTracker()
-        # self.__toBeFileWriteTracker = ToBeFileWriteTracker() #TODO
+        self.__toBeFileWriteTracker = ToBeFileWriteTracker()
 
         self.__progress_per_file: DefaultDict[TypedPath, SnapshotFileProgress] = (
             _keydefaultdict(lambda key: SnapshotFileProgress(self, key))  # type: ignore
@@ -247,7 +248,8 @@ def write_inline(self, literal_value: LiteralValue, call: CallStack):
     def write_to_be_file(
         self, path: TypedPath, data: "ByteString", call: CallStack
     ) -> None:
-        raise NotImplementedError
+        # Directly write to disk using ToBeFileWriteTracker
+        self.__toBeFileWriteTracker.writeToDisk(path, data, call, self._layout)
 
 
 class DiskStoragePytest(DiskStorage):

python/selfie-lib/selfie_lib/CacheSelfie.py

@@ -4,6 +4,7 @@
 from .Literals import LiteralValue, LiteralString, TodoStub
 from .SnapshotSystem import DiskStorage
 from .RoundTrip import Roundtrip
+import base64
 
 T = TypeVar("T")
 
@@ -75,3 +76,111 @@ def _to_be_impl(self, snapshot: Optional[str]) -> T:
                 raise Exception("Can't call `to_be_todo` in readonly mode!")
             else:
                 return self.roundtrip.parse(snapshot)
+
+
+class CacheSelfieBinary(Generic[T]):
+    def __init__(self, disk: DiskStorage, roundtrip: Roundtrip[T, bytes], generator):
+        self.disk: DiskStorage = disk
+        self.roundtrip: Roundtrip[T, bytes] = roundtrip
+        self.generator = generator
+
+    def to_match_disk(self, sub: str = "") -> T:
+        return self._to_match_disk_impl(sub, False)
+
+    def to_match_disk_TODO(self, sub: str = "") -> T:
+        return self._to_match_disk_impl(sub, True)
+
+    def _to_match_disk_impl(self, sub: str, is_todo: bool) -> T:
+        from .Selfie import _selfieSystem
+
+        system = _selfieSystem()
+        call = recordCall(False)
+
+        if system.mode.can_write(is_todo, call, system):
+            actual = self.generator()
+            serialized_data = self.roundtrip.serialize(actual)
+            self.disk.write_disk(Snapshot.of(serialized_data), sub, call)
+
+            if is_todo:
+                system.write_inline(TodoStub.to_match_disk.create_literal(), call)
+
+            return actual
+        else:
+            if is_todo:
+                raise Exception("Can't call `to_match_disk_TODO` in read-only mode!")
+            else:
+                snapshot = self.disk.read_disk(sub, call)
+
+                if snapshot is None:
+                    raise Exception(system.mode.msg_snapshot_not_found())
+
+                if snapshot.subject.is_binary or len(snapshot.facets) > 0:
+                    raise Exception(
+                        "Expected a binary subject with no facets, got {}".format(
+                            snapshot
+                        )
+                    )
+
+                return self.roundtrip.parse(snapshot.subject.value_binary())
+
+    def to_be_file_TODO(self, subpath: str) -> T:
+        return self._to_be_file_impl(subpath, True)
+
+    def to_be_file(self, subpath: str) -> T:
+        return self._to_be_file_impl(subpath, False)
+
+    def _to_be_file_impl(self, subpath: str, is_todo: bool) -> T:
+        from .Selfie import _selfieSystem
+
+        system = _selfieSystem()
+        call = recordCall(False)
+        writable = system.mode.can_write(is_todo, call, system)
+
+        if writable:
+            actual = self.generator()
+
+            if is_todo:
+                system.write_inline(TodoStub.to_be_file.create_literal(), call)
+
+            with open(subpath, "wb") as file:
+                file.write(self.roundtrip.serialize(actual))
+
+            return actual
+        else:
+            if is_todo:
+                raise Exception("Can't call `toBeFile_TODO` in read-only mode!")
+            else:
+                with open(subpath, "rb") as file:
+                    serialized_data = file.read()
+                return self.roundtrip.parse(serialized_data)
+
+    def to_be_base64_TODO(self, unused_arg: Optional[Any] = None) -> T:
+        return self._to_be_base64_impl(None)
+
+    def to_be_base64(self, snapshot: str) -> T:
+        return self._to_be_base64_impl(snapshot)
+
+    def _to_be_base64_impl(self, snapshot: Optional[str]) -> T:
+        from .Selfie import _selfieSystem
+
+        system = _selfieSystem()
+        call = recordCall(False)
+        writable = system.mode.can_write(snapshot is None, call, system)
+
+        if writable:
+            actual = self.generator()
+            base64_data = base64.b64encode(self.roundtrip.serialize(actual)).decode(
+                "utf-8"
+            )
+            literal_string_formatter = LiteralString()
+            system.write_inline(
+                LiteralValue(snapshot, base64_data, literal_string_formatter),
+                call,
+            )
+            return actual
+        else:
+            if snapshot is None:
+                raise Exception("Can't call `toBe_TODO` in read-only mode!")
+            else:
+                decoded_data = base64.b64decode(snapshot.encode("utf-8"))
+                return self.roundtrip.parse(decoded_data)

python/selfie-lib/selfie_lib/FS.py

@@ -33,5 +33,5 @@ def file_write_binary(self, typed_path: TypedPath, content: bytes):
         Path(typed_path.absolute_path).write_bytes(content)
 
     @abstractmethod
-    def assert_failed(self, message: str, expected=None, actual=None) -> Exception:
+    def assert_failed(self, message, expected=None, actual=None) -> Exception:
         pass

python/selfie-lib/selfie_lib/Selfie.py

@@ -1,4 +1,4 @@
-from typing import TypeVar, Any, Protocol, Union, overload
+from typing import Any, TypeVar, Optional, Protocol, Union, overload
 from .SelfieImplementations import ReprSelfie, StringSelfie
 from .SnapshotSystem import _selfieSystem
 from .Snapshot import Snapshot
@@ -39,15 +39,25 @@ def expect_selfie(
         return ReprSelfie(actual)
 
 
-def cache_selfie_string(to_cache: Cacheable[str]) -> CacheSelfie[str]:
-    """Create a CacheSelfie instance for caching strings with identity transformation."""
-    identity_roundtrip = Roundtrip.identity()
-    return cache_selfie_generic(identity_roundtrip, to_cache)
+@overload
+def cache_selfie(to_cache: Cacheable[str]) -> CacheSelfie[str]: ...
 
 
-def cache_selfie_generic(
-    roundtrip: Roundtrip[T, str], to_cache: Cacheable[T]
-) -> CacheSelfie[T]:
-    """Create a CacheSelfie instance for caching generic objects with specified roundtrip."""
-    deferred_disk_storage = _selfieSystem().disk_thread_local()
-    return CacheSelfie(deferred_disk_storage, roundtrip, to_cache)
+@overload
+def cache_selfie(
+    to_cache: Cacheable[T], roundtrip: Roundtrip[T, str]
+) -> CacheSelfie[T]: ...
+
+
+def cache_selfie(
+    to_cache: Union[Cacheable[str], Cacheable[T]],
+    roundtrip: Optional[Roundtrip[T, str]] = None,
+) -> Union[CacheSelfie[str], CacheSelfie[T]]:
+    if roundtrip is None:
+        # the cacheable better be a string!
+        return cache_selfie(to_cache, Roundtrip.identity())  # type: ignore
+    elif isinstance(roundtrip, Roundtrip) and to_cache is not None:
+        deferred_disk_storage = _selfieSystem().disk_thread_local()
+        return CacheSelfie(deferred_disk_storage, roundtrip, to_cache)
+    else:
+        raise TypeError("Invalid arguments provided to cache_selfie")

python/selfie-lib/selfie_lib/WriteTracker.py

@@ -239,3 +239,45 @@ def persist_writes(self, layout: SnapshotFileLayout):
 
         # Final write to disk for the last file processed
         layout.fs.file_write(current_file, content.as_string)
+
+
+class ToBeFileLazyBytes:
+    def __init__(self, location: TypedPath, layout: SnapshotFileLayout, data: bytes):
+        self.location = location
+        self.layout = layout
+        self.data = data
+
+    def writeToDisk(self) -> None:
+        if self.data is None:
+            raise Exception("Data has already been written to disk")
+        self.layout.fs.file_write_binary(self.location, self.data)
+        self.data = None  # Allow garbage collection
+
+    def readData(self):
+        if self.data is not None:
+            return self.data
+        return self.layout.fs.file_read_binary(self.location)
+
+    def __eq__(self, other):
+        if not isinstance(other, ToBeFileLazyBytes):
+            return False
+        return self.readData() == other.readData()
+
+    def __hash__(self):
+        return hash(self.readData())
+
+
+class ToBeFileWriteTracker(WriteTracker[TypedPath, ToBeFileLazyBytes]):
+    def __init__(self):
+        super().__init__()
+
+    def writeToDisk(
+        self,
+        key: TypedPath,
+        snapshot: bytes,
+        call: CallStack,
+        layout: SnapshotFileLayout,
+    ) -> None:
+        lazyBytes = ToBeFileLazyBytes(key, layout, snapshot)
+        self.recordInternal(key, lazyBytes, call, layout)
+        lazyBytes.writeToDisk()