feat: mpr #8 rilshok/state-impls

v.0.1.0 Basic implementation of main states

.github/workflows/publish.yml

@@ -25,6 +25,21 @@ jobs:
         name: python-package-distributions
         path: dist/
 
+  test:
+    name: Run tests with pytest
+    runs-on: ubuntu-latest
+    needs:
+    - build
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.10'
+    - name: Run tests
+      run: pytest
+
   publish-to-pypi:
     name: Publish Python 🐍 distribution 📦 to PyPI
     needs:

pyproject.toml

@@ -9,6 +9,7 @@ authors = [
     {name = "Vladislav A. Proskurov", email = "[email protected]"},
 ]
 dependencies = [
+    "jsonlines>=4.0.0",
     "humanize>=4.9.0",
     "typing-extensions>=4.8.0",
     "pytz>=2024.1",
@@ -17,6 +18,18 @@ dependencies = [
 [tool.setuptools.dynamic]
 version = {attr = "iokit.__version__"}
 
+[project.optional-dependencies]
+dev = ["iokit[lint,test]"]
+lint = [
+    "mypy",
+    "ruff",
+    "types-pytz",
+]
+test = [
+    "pytest",
+    "pytest-cov",
+]
+
 [project.urls]
 Homepage = "https://github.com/rilshok/iokit"
 Repository = "https://github.com/rilshok/iokit"

src/iokit/__init__.py

@@ -1,9 +1,18 @@
 __all__ = [
-    "State",
-    "Txt",
     "Gzip",
+    "Json",
+    "Jsonl",
+    "Tar",
+    "Txt",
+    "State",
+    "filter_states",
+    "find_state",
+    "load_file",
+    "save_file",
+    "save_temp",
 ]
-__version__ = "0.0.1"
+__version__ = "0.1.0"
 
-from .extensions import Gzip, Txt
-from .state import State
+from .extensions import Gzip, Json, Jsonl, Tar, Txt
+from .state import State, filter_states, find_state
+from .storage import load_file, save_file, save_temp

src/iokit/extensions/__init__.py

@@ -1,7 +1,13 @@
 __all__ = [
-    "Txt",
     "Gzip",
+    "Json",
+    "Jsonl",
+    "Tar",
+    "Txt",
 ]
 
 from .gz import Gzip
+from .json import Json
+from .jsonl import Jsonl
+from .tar import Tar
 from .txt import Txt

src/iokit/extensions/gz.py

@@ -6,7 +6,7 @@
 
 
 class Gzip(State, suffix="gz"):
-    def __init__(self, state: State, compression: int = 1, **kwargs: Any):
+    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:

src/iokit/extensions/json.py

@@ -0,0 +1,44 @@
+__all__ = [
+    "Json",
+]
+
+import json
+from functools import lru_cache
+from typing import Any, Callable
+
+from iokit.state import State
+
+
+@lru_cache
+def json_dumps(
+    *,
+    compact: bool,
+    ensure_ascii: bool,
+    allow_nan: bool,
+) -> Callable[[Any], str]:
+    item_sep = "," if compact else ", "
+    key_sep = ":" if compact else ": "
+    return json.JSONEncoder(
+        ensure_ascii=ensure_ascii,
+        allow_nan=allow_nan,
+        sort_keys=False,
+        separators=(item_sep, key_sep),
+    ).encode
+
+
+class Json(State, suffix="json"):
+    def __init__(
+        self,
+        data: Any,
+        *,
+        compact: bool = False,
+        ensure_ascii: bool = False,
+        allow_nan: bool = False,
+        **kwargs: Any,
+    ):
+        dumps = json_dumps(compact=compact, ensure_ascii=ensure_ascii, allow_nan=allow_nan)
+        data_ = dumps(data).encode("utf-8")
+        super().__init__(data=data_, **kwargs)
+
+    def load(self) -> Any:
+        return json.load(self.data)

src/iokit/extensions/jsonl.py

@@ -0,0 +1,34 @@
+__all__ = [
+    "Jsonl",
+]
+
+from io import BytesIO
+from typing import Any, Iterable
+
+from jsonlines import Reader, Writer
+
+from iokit.state import State
+
+from .json import json_dumps
+
+
+class Jsonl(State, suffix="jsonl"):
+    def __init__(
+        self,
+        sequence: Iterable[dict[str, Any]],
+        *,
+        compact: bool = True,
+        ensure_ascii: bool = False,
+        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)
+
+    def load(self) -> list[Any]:
+        with Reader(self.data) as reader:
+            return list(reader)

src/iokit/extensions/tar.py

@@ -0,0 +1,37 @@
+import tarfile
+from io import BytesIO
+from typing import Any, Iterable
+
+from iokit.state import State
+from iokit.tools.time import fromtimestamp
+
+
+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)
+
+        super().__init__(data=buffer, **kwargs)
+
+    def load(self) -> list[State]:
+        states: list[State] = []
+        with tarfile.open(fileobj=self.data, mode="r") as tar_buffer:
+            assert tar_buffer is not None
+            for member in tar_buffer.getmembers():
+                if not member.isfile():
+                    continue
+                member_buffer = tar_buffer.extractfile(member)
+                if member_buffer is None:
+                    continue
+                state = State(
+                    data=member_buffer.read(),
+                    name=member.name,
+                    time=fromtimestamp(member.mtime),
+                )
+                states.append(state)
+        return states

src/iokit/extensions/txt.py

@@ -5,7 +5,7 @@
 
 class Txt(State, suffix="txt"):
     def __init__(self, data: str, **kwargs: Any):
-        if not isinstance(data, str):  # type: ignore
+        if not isinstance(data, str):
             raise TypeError(f"Expected str, got {type(data).__name__}")
         super().__init__(data=data.encode("utf-8"), **kwargs)
 

src/iokit/py.typed

@@ -0,0 +1,4 @@
+This marker file declares that the package supports type checking.
+For details, you can refer to:
+- PEP561: https://www.python.org/dev/peps/pep-0561/
+- mypy docs: https://mypy.readthedocs.io/en/stable/installed_packages.html

src/iokit/state.py

@@ -1,16 +1,20 @@
+__all__ = [
+    "State",
+    "filter_states",
+    "find_state",
+]
+
 from datetime import datetime
+from fnmatch import fnmatch
 from io import BytesIO
-from typing import Any
+from typing import Any, Generator, Iterable
 
-import pytz
 from humanize import naturalsize
 from typing_extensions import Self
 
-Payload = BytesIO | bytes
-
+from iokit.tools.time import now
 
-def now() -> datetime:
-    return datetime.now(pytz.utc)
+Payload = BytesIO | bytes
 
 
 class StateName:
@@ -45,6 +49,14 @@ def __str__(self) -> str:
     def __repr__(self) -> str:
         return self._name
 
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, str):
+            return self._name == other
+        if isinstance(other, StateName):
+            return self._name == other._name
+        msg = f"Expected str or StateName, got {type(other).__name__}"
+        raise NotImplementedError(msg)
+
     @classmethod
     def make(cls, stem: "str | StateName", suffix: str) -> Self:
         if suffix:
@@ -131,3 +143,15 @@ def load(self) -> Any:
         if not self.name.suffix:
             return self.data.getvalue()
         return self.cast().load()
+
+
+def filter_states(states: Iterable[State], pattern: str) -> Generator[State, None, None]:
+    for state in states:
+        if fnmatch(str(state.name), pattern):
+            yield state
+
+
+def find_state(states: Iterable[State], pattern: str) -> State:
+    for state in filter_states(states, pattern):
+        return state
+    raise FileNotFoundError(f"State not found: {pattern!r}")

src/iokit/storage/__init__.py

@@ -0,0 +1,7 @@
+__all__ = [
+    "load_file",
+    "save_file",
+    "save_temp",
+]
+
+from .local import load_file, save_file, save_temp

src/iokit/storage/local.py

@@ -0,0 +1,46 @@
+__all__ = [
+    "load_file",
+    "save_file",
+    "save_temp",
+]
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Generator
+
+from iokit.state import State
+from iokit.tools.time import fromtimestamp
+
+PathLike = str | Path
+
+
+def load_file(path: PathLike) -> State:
+    path = Path(path).resolve()
+    mtime = fromtimestamp(path.stat().st_mtime)
+    return State(data=path.read_bytes(), name=path.name, time=mtime).cast()
+
+
+def save_file(
+    state: State,
+    root: PathLike = "",
+    parents: bool = False,
+    force: bool = False,
+) -> Path:
+    root = Path(root).resolve()
+    path = (root / str(state.name)).resolve()
+    if not path.is_relative_to(root):
+        msg = f"Path is outside of root: root='{root!s}', state.name='{state.name!s}'"
+        raise ValueError(msg)
+    if path.exists() and not force:
+        msg = f"File already exists: path='{path!s}'"
+        raise FileExistsError(msg)
+    root.mkdir(parents=parents, exist_ok=True)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_bytes(state.data.getvalue())
+    return path
+
+
+@contextmanager
+def save_temp(state: State) -> Generator[Path, None, None]:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        yield save_file(state, root=temp_dir)

src/iokit/tools/time.py

@@ -0,0 +1,11 @@
+from datetime import datetime
+
+from pytz import utc
+
+
+def fromtimestamp(timestamp: float) -> datetime:
+    return datetime.fromtimestamp(timestamp, utc)
+
+
+def now() -> datetime:
+    return datetime.now(utc)

tests/test_filter_states.py

@@ -0,0 +1,31 @@
+from typing import Iterable
+
+from iokit import Json, State, filter_states
+
+
+def filter_states_(states: Iterable[State], pattern: str) -> list[State]:
+    return list(filter_states(states, pattern))
+
+
+def test_filter_states() -> None:
+    banana = Json({"name": "banana"}, name="banana")
+    tomato = Json({"name": "tomato"}, name="tomato")
+    orange = Json({"name": "orange"}, name="orange")
+    cherry = Json({"name": "cherry"}, name="cherry")
+    potato = Json({"name": "potato"}, name="potato")
+
+    states = [banana, tomato, orange, cherry, potato]
+
+    assert filter_states_(states, "") == []
+    assert filter_states_(states, "*") == states
+    assert filter_states_(states, "o*") == [orange]
+    assert filter_states_(states, "o*") == [orange]
+    assert filter_states_(states, "x*") == []
+    assert filter_states_(states, "b*n") == [banana]
+    assert filter_states_(states, "c*") == [cherry]
+    assert filter_states_(states, "b*n*") == [banana]
+    assert filter_states_(states, "p*t*") == [potato]
+    assert filter_states_(states, "b*n*o") == []
+    assert filter_states_(states, "[bpt]*") == [banana, tomato, potato]
+    assert filter_states_(states, "[*") == []
+    assert filter_states_(states, "t?mato*") == [tomato]