CacheSelfie.py

Go to the documentation of this file.

import base64
from typing import Any, Callable, Generic, Optional, TypeVar, Union, overload
 
from .Literals import LiteralString, LiteralValue, TodoStub
from .Roundtrip import Roundtrip
from .Snapshot import Snapshot
from .SnapshotSystem import DiskStorage, _selfieSystem
from .WriteTracker import recordCall
 
T = TypeVar("T", covariant=True)
 
 
@overload
def cache_selfie(to_cache: Callable[..., str]) -> "CacheSelfie[str]": ...
 
 
@overload
def cache_selfie(
    to_cache: Callable[..., T], roundtrip: Roundtrip[T, str]
) -> "CacheSelfie[T]": ...
 
 
def cache_selfie(
    to_cache: Union[Callable[..., str], Callable[..., 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)  # type: ignore
    else:
        raise TypeError("Invalid arguments provided to cache_selfie")
 
 
def cache_selfie_json(to_cache: Callable[..., T]) -> "CacheSelfie[T]":
    return cache_selfie(to_cache, Roundtrip.json())
 
 
@overload
def cache_selfie_binary(
    to_cache: Callable[..., bytes],
) -> "CacheSelfieBinary[bytes]": ...
 
 
@overload
def cache_selfie_binary(
    to_cache: Callable[..., T], roundtrip: Roundtrip[T, bytes]
) -> "CacheSelfieBinary[T]": ...
 
 
def cache_selfie_binary(
    to_cache: Union[Callable[..., bytes], Callable[..., T]],
    roundtrip: Optional[Roundtrip[T, bytes]] = None,
) -> Union["CacheSelfieBinary[bytes]", "CacheSelfieBinary[T]"]:
    if roundtrip is None:
        # the cacheable better be a bytes!
        return cache_selfie_binary(to_cache, Roundtrip.identity())  # type: ignore
    elif isinstance(roundtrip, Roundtrip) and to_cache is not None:
        deferred_disk_storage = _selfieSystem().disk_thread_local()
        return CacheSelfieBinary(deferred_disk_storage, roundtrip, to_cache)  # type: ignore
    else:
        raise TypeError("Invalid arguments provided to cache_selfie")
 
 
class CacheSelfie(Generic[T]):
    def __init__(
        self,
        disk: DiskStorage,
        roundtrip: Roundtrip[T, str],
        generator: Callable[..., T],
    ):
        self.disk = disk
        self.roundtrip = 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:
        call = recordCall(False)
        system = _selfieSystem()
        if system.mode.can_write(is_todo, call, system):
            actual = self.generator()
            self.disk.write_disk(
                Snapshot.of(self.roundtrip.serialize(actual)), 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 readonly 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(
                        f"Expected a string subject with no facets, got {snapshot}"
                    )
                return self.roundtrip.parse(snapshot.subject.value_string())
 
    def to_be_TODO(self, _: Optional[Any] = None) -> T:
        return self._to_be_impl(None)
 
    def to_be(self, expected: str) -> T:
        return self._to_be_impl(expected)
 
    def _to_be_impl(self, snapshot: Optional[str]) -> T:
        call = recordCall(False)
        system = _selfieSystem()
        writable = system.mode.can_write(snapshot is None, call, system)
        if writable:
            actual = self.generator()
            literal_string_formatter = LiteralString()
            system.write_inline(
                LiteralValue(
                    snapshot, self.roundtrip.serialize(actual), literal_string_formatter
                ),
                call,
            )
            return actual
        else:
            if snapshot is None:
                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: Callable[..., T],
    ):
        self.disk = disk
        self.roundtrip = 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:
        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(
                        f"Expected a binary subject with no facets, got {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:
        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 `to_be_file_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, _: 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:
        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 `to_be_TODO` in read-only mode!")
            else:
                decoded_data = base64.b64decode(snapshot.encode("utf-8"))
                return self.roundtrip.parse(decoded_data)