diff --git a/src/iokit/__init__.py b/src/iokit/__init__.py index 14ded44..f3a658b 100644 --- a/src/iokit/__init__.py +++ b/src/iokit/__init__.py @@ -1,33 +1,36 @@ +__version__ = "0.2.0" __all__ = [ "Dat", - "Encryption", + "decrypt", + "download_file", + "Enc", + "encrypt", "Env", + "filter_states", + "find_state", + "Flac", "Gzip", "Json", "Jsonl", - "Tar", - "Txt", - "Yaml", - "Zip", + "load_file", "Mp3", "Npy", - "Wav", - "Waveform", - "Flac", "Ogg", - "State", - "filter_states", - "find_state", - "download_file", - "load_file", "save_file", "save_temp", + "SecretState", + "State", + "Tar", + "Txt", + "Waveform", + "Wav", + "Yaml", + "Zip", ] -__version__ = "0.2.0" from .extensions import ( Dat, - Encryption, + Enc, Env, Flac, Gzip, @@ -36,12 +39,15 @@ Mp3, Npy, Ogg, + SecretState, Tar, Txt, Wav, Waveform, Yaml, Zip, + decrypt, + encrypt, ) from .state import State, filter_states, find_state from .storage import download_file, load_file, save_file, save_temp diff --git a/src/iokit/extensions/__init__.py b/src/iokit/extensions/__init__.py index 0ea2ed8..9f7950a 100644 --- a/src/iokit/extensions/__init__.py +++ b/src/iokit/extensions/__init__.py @@ -1,25 +1,28 @@ __all__ = [ - "Flac", - "Mp3", - "Ogg", - "Wav", - "Waveform", "Dat", - "Encryption", + "Enc", "Env", + "Flac", "Gzip", "Json", "Jsonl", + "Mp3", "Npy", + "Ogg", + "SecretState", "Tar", "Txt", + "Wav", + "Waveform", "Yaml", "Zip", + "decrypt", + "encrypt", ] from .audio import Flac, Mp3, Ogg, Wav, Waveform from .dat import Dat -from .enc import Encryption +from .enc import Enc, SecretState, decrypt, encrypt from .env import Env from .gz import Gzip from .json import Json diff --git a/src/iokit/extensions/audio.py b/src/iokit/extensions/audio.py index 7bffb78..7f39391 100644 --- a/src/iokit/extensions/audio.py +++ b/src/iokit/extensions/audio.py @@ -1,3 +1,5 @@ +__all__ = ["Waveform", "Flac", "Wav", "Mp3", "Ogg"] + from dataclasses import dataclass from io import BytesIO from typing import Any @@ -11,16 +13,17 @@ class AudioState(State, suffix=""): def __init__(self, waveform: "Waveform", **kwargs: Any): - soundfile.write( - file=(target := BytesIO()), - data=waveform.wave, - samplerate=waveform.freq, - format=self._suffix, - ) - super().__init__(data=target.getvalue(), **kwargs) + with BytesIO() as buffer: + soundfile.write( + file=buffer, + data=waveform.wave, + samplerate=waveform.freq, + format=self._suffix, + ) + super().__init__(data=buffer.getvalue(), **kwargs) def load(self) -> "Waveform": - wave, freq = soundfile.read(self.data, always_2d=True) + wave, freq = soundfile.read(self.buffer, always_2d=True) return Waveform(wave=wave, freq=freq) diff --git a/src/iokit/extensions/dat.py b/src/iokit/extensions/dat.py index 0cc3b02..e474256 100644 --- a/src/iokit/extensions/dat.py +++ b/src/iokit/extensions/dat.py @@ -1,15 +1,13 @@ -__all__ = [ - "Dat", -] +__all__ = ["Dat"] from typing import Any -from iokit.state import Payload, State +from iokit.state import State class Dat(State, suffix="dat"): - def __init__(self, data: Payload, **kwargs: Any): + def __init__(self, data: bytes, **kwargs: Any): super().__init__(data=data, **kwargs) def load(self) -> bytes: - return self._data.getvalue() + return self.data diff --git a/src/iokit/extensions/enc.py b/src/iokit/extensions/enc.py index 5648f49..aa32633 100644 --- a/src/iokit/extensions/enc.py +++ b/src/iokit/extensions/enc.py @@ -1,3 +1,5 @@ +__all__ = ["SecretState", "Enc", "encrypt", "decrypt"] + import struct from collections.abc import Iterator from hashlib import sha256 @@ -91,12 +93,12 @@ def __repr__(self) -> str: @classmethod def pack(cls, state: State, password: bytes | str, salt: bytes | str = b"42") -> Self: - payload = _pack_arrays(str(state.name).encode("utf-8"), state.data.getvalue()) + payload = _pack_arrays(str(state.name).encode("utf-8"), state.data) data = encrypt(data=payload, password=_to_bytes(password), salt=_to_bytes(salt)) return cls(data=data) -class Encryption(State, suffix="enc"): +class Enc(State, suffix="enc"): def __init__( self, state: State, @@ -112,4 +114,4 @@ def __init__( super().__init__(data=data, name=name, **kwargs) def load(self) -> SecretState: - return SecretState(data=self.data.getvalue()) + return SecretState(data=self.data) diff --git a/src/iokit/extensions/env.py b/src/iokit/extensions/env.py index d0595f3..7a75feb 100644 --- a/src/iokit/extensions/env.py +++ b/src/iokit/extensions/env.py @@ -1,3 +1,5 @@ +__all__ = ["Env"] + from io import StringIO from pathlib import Path from tempfile import TemporaryDirectory @@ -20,5 +22,5 @@ def __init__(self, data: dict[str, str], **kwargs: Any): super().__init__(data=data_bytes, **kwargs) def load(self) -> dict[str, str | None]: - stream = StringIO(self.data.getvalue().decode()) - return dict(dotenv.dotenv_values(stream=stream)) + with StringIO(self.data.decode()) as stream: + return dict(dotenv.dotenv_values(stream=stream)) diff --git a/src/iokit/extensions/gz.py b/src/iokit/extensions/gz.py index 72f5b0f..9891fb6 100644 --- a/src/iokit/extensions/gz.py +++ b/src/iokit/extensions/gz.py @@ -1,3 +1,5 @@ +__all__ = ["Gzip"] + import gzip from io import BytesIO from typing import Any @@ -7,14 +9,13 @@ class Gzip(State, suffix="gz"): def __init__(self, state: State, *, compression: int = 1, **kwargs: Any): - data = BytesIO() - gzip_file = gzip.GzipFile(fileobj=data, mode="wb", compresslevel=compression, mtime=0) - with gzip_file as gzip_buffer: - gzip_buffer.write(state.data.getvalue()) - super().__init__(data=data, name=state.name, **kwargs) + with BytesIO() as buffer: + gzip_file = gzip.GzipFile(fileobj=buffer, mode="wb", compresslevel=compression, mtime=0) + with gzip_file as gzip_buffer: + gzip_buffer.write(state.data) + super().__init__(data=buffer.getvalue(), name=state.name, **kwargs) def load(self) -> State: - gzip_file = gzip.GzipFile(fileobj=self.data, mode="rb") - with gzip_file as file: - data = file.read() - return State(data=data, name=str(self.name).removesuffix(".gz")).cast() + # gzip_file = + with gzip.GzipFile(fileobj=self.buffer, mode="rb") as file: + return State(data=file.read(), name=str(self.name).removesuffix(".gz")).cast() diff --git a/src/iokit/extensions/json.py b/src/iokit/extensions/json.py index 3efe604..5ebfa89 100644 --- a/src/iokit/extensions/json.py +++ b/src/iokit/extensions/json.py @@ -1,6 +1,4 @@ -__all__ = [ - "Json", -] +__all__ = ["Json"] import json from collections.abc import Callable @@ -42,4 +40,4 @@ def __init__( super().__init__(data=data_, **kwargs) def load(self) -> Any: - return json.load(self.data) + return json.load(self.buffer) diff --git a/src/iokit/extensions/jsonl.py b/src/iokit/extensions/jsonl.py index c176283..6fdaedc 100644 --- a/src/iokit/extensions/jsonl.py +++ b/src/iokit/extensions/jsonl.py @@ -1,6 +1,4 @@ -__all__ = [ - "Jsonl", -] +__all__ = ["Jsonl"] from collections.abc import Iterable from io import BytesIO @@ -23,13 +21,13 @@ def __init__( allow_nan: bool = False, **kwargs: Any, ): - buffer = BytesIO() - dumps = json_dumps(compact=compact, ensure_ascii=ensure_ascii, allow_nan=allow_nan) - with Writer(buffer, compact=compact, sort_keys=False, dumps=dumps) as writer: - for item in sequence: - writer.write(item) - super().__init__(data=buffer, **kwargs) + with BytesIO() as buffer: + dumps = json_dumps(compact=compact, ensure_ascii=ensure_ascii, allow_nan=allow_nan) + with Writer(buffer, compact=compact, sort_keys=False, dumps=dumps) as writer: + for item in sequence: + writer.write(item) + super().__init__(data=buffer.getvalue(), **kwargs) def load(self) -> list[Any]: - with Reader(self.data) as reader: + with Reader(self.buffer) as reader: return list(reader) diff --git a/src/iokit/extensions/npy.py b/src/iokit/extensions/npy.py index 4377265..a31a02c 100644 --- a/src/iokit/extensions/npy.py +++ b/src/iokit/extensions/npy.py @@ -1,3 +1,5 @@ +__all__ = ["Npy"] + from io import BytesIO from typing import Any @@ -14,4 +16,4 @@ def __init__(self, array: NDArray[Any], **kwargs: Any) -> None: super().__init__(data=buffer.getvalue(), **kwargs) def load(self) -> NDArray[Any]: - return np.load(self.data, allow_pickle=False, fix_imports=False) + return np.load(self.buffer, allow_pickle=False, fix_imports=False) diff --git a/src/iokit/extensions/tar.py b/src/iokit/extensions/tar.py index adacdf7..ab44c3c 100644 --- a/src/iokit/extensions/tar.py +++ b/src/iokit/extensions/tar.py @@ -1,3 +1,5 @@ +__all__ = ["Tar"] + import tarfile from collections.abc import Iterable from io import BytesIO @@ -9,19 +11,19 @@ class Tar(State, suffix="tar"): def __init__(self, states: Iterable[State], **kwargs: Any): - buffer = BytesIO() - with tarfile.open(fileobj=buffer, mode="w") as tar_buffer: - for state in states: - file_data = tarfile.TarInfo(name=str(state.name)) - file_data.size = state.size - file_data.mtime = int(state.time.timestamp()) - tar_buffer.addfile(fileobj=state.data, tarinfo=file_data) + with BytesIO() as buffer: + with tarfile.open(fileobj=buffer, mode="w") as tar_buffer: + for state in states: + file_data = tarfile.TarInfo(name=str(state.name)) + file_data.size = state.size + file_data.mtime = int(state.time.timestamp()) + tar_buffer.addfile(fileobj=state.buffer, tarinfo=file_data) - super().__init__(data=buffer, **kwargs) + super().__init__(data=buffer.getvalue(), **kwargs) def load(self) -> list[State]: states: list[State] = [] - with tarfile.open(fileobj=self.data, mode="r") as tar_buffer: + with tarfile.open(fileobj=self.buffer, mode="r") as tar_buffer: assert tar_buffer is not None for member in tar_buffer.getmembers(): if not member.isfile(): diff --git a/src/iokit/extensions/txt.py b/src/iokit/extensions/txt.py index bf3a22b..bd035f6 100644 --- a/src/iokit/extensions/txt.py +++ b/src/iokit/extensions/txt.py @@ -1,3 +1,5 @@ +__all__ = ["Txt"] + from typing import Any from iokit.state import State @@ -5,9 +7,7 @@ class Txt(State, suffix="txt"): def __init__(self, data: str, **kwargs: Any): - if not isinstance(data, str): - raise TypeError(f"Expected str, got {type(data).__name__}") super().__init__(data=data.encode("utf-8"), **kwargs) def load(self) -> str: - return self.data.getvalue().decode("utf-8") + return self.data.decode("utf-8") diff --git a/src/iokit/extensions/yaml.py b/src/iokit/extensions/yaml.py index 493013d..a053f42 100644 --- a/src/iokit/extensions/yaml.py +++ b/src/iokit/extensions/yaml.py @@ -1,6 +1,4 @@ -__all__ = [ - "Yaml", -] +__all__ = ["Yaml"] from typing import Any @@ -15,4 +13,4 @@ def __init__(self, data: Any, **kwargs: Any): super().__init__(data=data, **kwargs) def load(self) -> Any: - return yaml.safe_load(self.data) + return yaml.safe_load(self.buffer) diff --git a/src/iokit/extensions/zip.py b/src/iokit/extensions/zip.py index ea69071..8875895 100644 --- a/src/iokit/extensions/zip.py +++ b/src/iokit/extensions/zip.py @@ -1,3 +1,5 @@ +__all__ = ["Zip"] + import zipfile from collections.abc import Iterable from datetime import datetime @@ -9,16 +11,16 @@ class Zip(State, suffix="zip"): def __init__(self, states: Iterable[State], **kwargs: Any): - buffer = BytesIO() - with zipfile.ZipFile(buffer, mode="w") as zip_buffer: - for state in states: - zip_buffer.writestr(str(state.name), data=state.data.getvalue()) + with BytesIO() as buffer: + with zipfile.ZipFile(buffer, mode="w") as zip_buffer: + for state in states: + zip_buffer.writestr(str(state.name), data=state.data) - super().__init__(data=buffer, **kwargs) + super().__init__(data=buffer.getvalue(), **kwargs) def load(self) -> list[State]: states: list[State] = [] - with zipfile.ZipFile(self.data, mode="r") as zip_buffer: + with zipfile.ZipFile(self.buffer, mode="r") as zip_buffer: for file in zip_buffer.namelist(): with zip_buffer.open(file) as member_buffer: state = State( diff --git a/src/iokit/state.py b/src/iokit/state.py index 48d5081..b8ffadd 100644 --- a/src/iokit/state.py +++ b/src/iokit/state.py @@ -16,8 +16,6 @@ from iokit.tools.time import now -Payload = BytesIO | bytes - class StateName: def __init__(self, name: str): @@ -70,8 +68,8 @@ class State: _suffix: str = "" _suffixes: tuple[str, ...] = ("",) - def __init__(self, data: Payload, name: str | StateName = "", time: datetime | None = None): - self._data = BytesIO(data) if isinstance(data, bytes) else data + def __init__(self, data: bytes, name: str | StateName = "", time: datetime | None = None): + self._data = data self._name = StateName.make(name, self._suffix) self._time = time or now() @@ -115,13 +113,16 @@ def time(self, value: datetime) -> None: self._time = value @property - def data(self) -> BytesIO: - self._data.seek(0) - return BytesIO(self._data.getvalue()) + def data(self) -> bytes: + return self._data + + @property + def buffer(self) -> BytesIO: + return BytesIO(self._data) @property def size(self) -> int: - return self._data.getbuffer().nbytes + return len(self._data) def __repr__(self) -> str: size = naturalsize(self.size, gnu=True) @@ -148,7 +149,7 @@ def cast(self) -> "State": def load(self) -> Any: if not self.name.suffix: - return self.data.getvalue() + return self.data state = self.cast() if type(state) is State: # pylint: disable=unidiomatic-typecheck msg = f"Cannot load state with suffix '{self.name.suffix}'" diff --git a/src/iokit/storage/local.py b/src/iokit/storage/local.py index c5274ae..f0b1b01 100644 --- a/src/iokit/storage/local.py +++ b/src/iokit/storage/local.py @@ -36,7 +36,7 @@ def save_file( raise FileExistsError(msg) root.mkdir(parents=parents, exist_ok=True) path.parent.mkdir(parents=True, exist_ok=True) - path.write_bytes(state.data.getvalue()) + path.write_bytes(state.data) return path diff --git a/tests/test_download_file.py b/tests/test_download_file.py index 615a5c3..703bba6 100644 --- a/tests/test_download_file.py +++ b/tests/test_download_file.py @@ -4,4 +4,4 @@ def test_download_file() -> None: uri = "https://raw.githubusercontent.com/rilshok/iokit/main/LICENSE" state = download_file(uri) - assert "MIT License" in state.data.getvalue().decode("utf-8") + assert "MIT License" in state.data.decode("utf-8") diff --git a/tests/test_encryption.py b/tests/test_encryption.py index 4da3b96..abc92d9 100644 --- a/tests/test_encryption.py +++ b/tests/test_encryption.py @@ -2,7 +2,7 @@ import pytest -from iokit import Encryption, Json +from iokit import Enc, Json def test_encryption() -> None: @@ -14,7 +14,7 @@ def test_encryption() -> None: "int": 42, } json = Json(data, name="different") - state = Encryption(json, password="pA$sw0Rd", salt="s@lt") + state = Enc(json, password="pA$sw0Rd", salt="s@lt") state_secret = state.load() with pytest.raises(ValueError) as excinfo: state_secret.load(password="pA$sw0Rd") diff --git a/tests/test_json.py b/tests/test_json.py index 4dba62e..45c24ee 100644 --- a/tests/test_json.py +++ b/tests/test_json.py @@ -7,21 +7,21 @@ def test_json_empty() -> None: state = Json({}, name="empty") assert state.name == "empty.json" assert state.size == 2 - assert state.data.getvalue() == b"{}" + assert state.data == b"{}" assert not state.load() def test_json_single() -> None: state = Json({"key": "value"}, name="single") assert state.name == "single.json" - assert state.data.getvalue() == b'{"key": "value"}' + assert state.data == b'{"key": "value"}' assert state.load() == {"key": "value"} def test_json_multiple() -> None: state = Json({"first": 1, "second": 2}, name="multiple") assert state.name == "multiple.json" - assert state.data.getvalue() == b'{"first": 1, "second": 2}' + assert state.data == b'{"first": 1, "second": 2}' assert state.load() == {"first": 1, "second": 2} assert 20 < state.size < 30 @@ -48,11 +48,11 @@ def test_json_is_string() -> None: state = Json("hello", name="string") assert state.load() == "hello" assert state.size == 7 - assert state.data.getvalue() == b'"hello"' + assert state.data == b'"hello"' def test_json_is_sequence() -> None: state = Json([1, 2, 3], name="sequence") assert state.load() == [1, 2, 3] assert state.size == 9 - assert state.data.getvalue() == b"[1, 2, 3]" + assert state.data == b"[1, 2, 3]" diff --git a/tests/test_jsonl.py b/tests/test_jsonl.py index a308582..782ba99 100644 --- a/tests/test_jsonl.py +++ b/tests/test_jsonl.py @@ -5,21 +5,21 @@ def test_jsonl_empty() -> None: state = Jsonl([], name="empty") assert state.name == "empty.jsonl" assert state.size == 0 - assert state.data.getvalue() == b"" + assert state.data == b"" assert not state.load() def test_jsonl_single() -> None: state = Jsonl([{"key": "value"}], name="single") assert state.name == "single.jsonl" - assert state.data.getvalue() == b'{"key":"value"}\n' + assert state.data == b'{"key":"value"}\n' assert state.load() == [{"key": "value"}] def test_jsonl_multiple_repeat() -> None: state = Jsonl([{"key": "value"}] * 2, name="multiple") assert state.name == "multiple.jsonl" - assert state.data.getvalue() == b'{"key":"value"}\n{"key":"value"}\n' + assert state.data == b'{"key":"value"}\n{"key":"value"}\n' assert state.load() == [{"key": "value"}] * 2 diff --git a/tests/test_tar.py b/tests/test_tar.py index 0577949..b53e884 100644 --- a/tests/test_tar.py +++ b/tests/test_tar.py @@ -22,7 +22,7 @@ def test_tar_compress() -> None: archive2_gz = Gzip(archive) assert archive1_gz.size == archive2_gz.size - assert archive1_gz.data.getvalue() == archive2_gz.data.getvalue() + assert archive1_gz.data == archive2_gz.data loaded = archive1_gz.load().load() assert find_state(loaded, "text1.txt").load() == "First file" assert find_state(loaded, "text2.txt").load() == "Second file" diff --git a/tests/test_yaml.py b/tests/test_yaml.py index 069ad5f..42fc1e3 100644 --- a/tests/test_yaml.py +++ b/tests/test_yaml.py @@ -5,7 +5,7 @@ def test_yaml_empty() -> None: state = Yaml([], name="empty") assert state.name == "empty.yaml" assert state.size == 3 - assert state.data.getvalue() == b"[]\n" + assert state.data == b"[]\n" assert not state.load()