omlish 0.0.0.dev133__py3-none-any.whl → 0.0.0.dev177__py3-none-any.whl

omlish/io/compress/zlib.py

@@ -0,0 +1,74 @@
+import dataclasses as dc
+import functools
+import typing as ta
+
+from ... import lang
+from ..generators import BytesSteppedGenerator
+from ..generators import BytesSteppedReaderGenerator
+from .adapters import CompressorObjectIncrementalAdapter
+from .adapters import DecompressorObjectIncrementalAdapter
+from .base import Compression
+from .base import IncrementalCompression
+from .codecs import make_compression_codec
+from .codecs import make_compression_lazy_loaded_codec
+
+
+if ta.TYPE_CHECKING:
+    import zlib
+else:
+    zlib = lang.proxy_import('zlib')
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class ZlibCompression(Compression, IncrementalCompression):
+    level: int = 9
+
+    wbits: int | None = None
+    strategy: int | None = None
+    zdict: bytes | None = None
+
+    def compress(self, d: bytes) -> bytes:
+        return zlib.compress(
+            d,
+            self.level,
+            **(dict(wbits=self.wbits) if self.wbits is not None else {}),
+        )
+
+    def decompress(self, d: bytes) -> bytes:
+        return zlib.decompress(
+            d,
+            **(dict(wbits=self.wbits) if self.wbits is not None else {}),
+        )
+
+    def compress_incremental(self) -> BytesSteppedGenerator[None]:
+        return lang.nextgen(CompressorObjectIncrementalAdapter(
+            functools.partial(
+                zlib.compressobj,  # type: ignore
+                self.level,
+                **(dict(wbits=self.wbits) if self.wbits is not None else {}),  # type: ignore[arg-type]
+                **(dict(strategy=self.strategy) if self.strategy is not None else {}),  # type: ignore[arg-type]
+                **(dict(zdict=self.zdict) if self.zdict is not None else {}),  # type: ignore[arg-type]
+            ),
+        )())
+
+    def decompress_incremental(self) -> BytesSteppedReaderGenerator[None]:
+        return DecompressorObjectIncrementalAdapter(
+            functools.partial(  # type: ignore
+                zlib.decompressobj,
+                **(dict(wbits=self.wbits) if self.wbits is not None else {}),  # type: ignore[arg-type]
+                **(dict(zdict=self.zdict) if self.zdict is not None else {}),  # type: ignore[arg-type]
+            ),
+            trailing_error=OSError,
+        )()
+
+
+##
+
+
+ZLIB_CODEC = make_compression_codec('zlib', ZlibCompression)
+
+# @omlish-manifest
+_ZLIB_LAZY_CODEC = make_compression_lazy_loaded_codec(__name__, 'ZLIB_CODEC', ZLIB_CODEC)

omlish/io/compress/zstd.py

@@ -0,0 +1,44 @@
+import dataclasses as dc
+import typing as ta
+
+from ... import lang
+from .base import Compression
+from .codecs import make_compression_codec
+from .codecs import make_compression_lazy_loaded_codec
+
+
+if ta.TYPE_CHECKING:
+    import zstandard
+else:
+    zstandard = lang.proxy_import('zstandard')
+
+
+##
+
+
+@dc.dataclass(frozen=True, kw_only=True)
+class ZstdCompression(Compression):
+    level: int | None = None
+
+    max_output_size: int = 0
+
+    def compress(self, d: bytes) -> bytes:
+        return zstandard.compress(
+            d,
+            **(dict(level=self.level) if self.level is not None else {}),
+        )
+
+    def decompress(self, d: bytes) -> bytes:
+        return zstandard.decompress(
+            d,
+            max_output_size=self.max_output_size,
+        )
+
+
+##
+
+
+ZSTD_CODEC = make_compression_codec('zstd', ZstdCompression)
+
+# @omlish-manifest
+_ZSTD_LAZY_CODEC = make_compression_lazy_loaded_codec(__name__, 'ZSTD_CODEC', ZSTD_CODEC)

omlish/io/fdio/__init__.py

@@ -0,0 +1 @@
+# @omlish-lite

omlish/{lite → io}/fdio/handlers.py

@@ -3,11 +3,11 @@ import abc
 import socket
 import typing as ta
 
-from ..check import check_not_none
-from ..socket import SocketAddress
+from ...lite.check import check
+from ...lite.socket import SocketAddress
 
 
-class FdIoHandler(abc.ABC):
+class FdioHandler(abc.ABC):
     @abc.abstractmethod
     def fd(self) -> int:
         raise NotImplementedError
@@ -43,7 +43,7 @@ class FdIoHandler(abc.ABC):
         pass
 
 
-class SocketFdIoHandler(FdIoHandler, abc.ABC):
+class SocketFdioHandler(FdioHandler, abc.ABC):
     def __init__(
             self,
             addr: SocketAddress,
@@ -55,7 +55,7 @@ class SocketFdIoHandler(FdIoHandler, abc.ABC):
         self._sock: ta.Optional[socket.socket] = sock
 
     def fd(self) -> int:
-        return check_not_none(self._sock).fileno()
+        return check.not_none(self._sock).fileno()
 
     @property
     def closed(self) -> bool:

omlish/{lite → io}/fdio/kqueue.py

@@ -4,13 +4,13 @@ import select
 import sys
 import typing as ta
 
-from .pollers import FdIoPoller
+from .pollers import FdioPoller
 
 
-KqueueFdIoPoller: ta.Optional[ta.Type[FdIoPoller]]
+KqueueFdioPoller: ta.Optional[ta.Type[FdioPoller]]
 if sys.platform == 'darwin' or sys.platform.startswith('freebsd'):
 
-    class _KqueueFdIoPoller(FdIoPoller):
+    class _KqueueFdioPoller(FdioPoller):
         DEFAULT_MAX_EVENTS = 1000
 
         def __init__(
@@ -98,14 +98,14 @@ if sys.platform == 'darwin' or sys.platform.startswith('freebsd'):
 
         #
 
-        def poll(self, timeout: ta.Optional[float]) -> FdIoPoller.PollResult:
+        def poll(self, timeout: ta.Optional[float]) -> FdioPoller.PollResult:
             kq = self._get_kqueue()
             try:
                 kes = kq.control(None, self._max_events, timeout)
 
             except OSError as exc:
                 if exc.errno == errno.EINTR:
-                    return FdIoPoller.PollResult(msg='EINTR encountered in poll', exc=exc)
+                    return FdioPoller.PollResult(msg='EINTR encountered in poll', exc=exc)
                 else:
                     raise
 
@@ -117,8 +117,8 @@ if sys.platform == 'darwin' or sys.platform.startswith('freebsd'):
                 if ke.filter == select.KQ_FILTER_WRITE:
                     w.append(ke.ident)
 
-            return FdIoPoller.PollResult(r, w)
+            return FdioPoller.PollResult(r, w)
 
-    KqueueFdIoPoller = _KqueueFdIoPoller
+    KqueueFdioPoller = _KqueueFdioPoller
 else:
-    KqueueFdIoPoller = None
+    KqueueFdioPoller = None

omlish/{lite → io}/fdio/manager.py

@@ -1,27 +1,27 @@
 # ruff: noqa: UP006 UP007
 import typing as ta
 
-from .handlers import FdIoHandler
-from .pollers import FdIoPoller
+from .handlers import FdioHandler
+from .pollers import FdioPoller
 
 
-class FdIoManager:
+class FdioManager:
     def __init__(
             self,
-            poller: FdIoPoller,
+            poller: FdioPoller,
     ) -> None:
         super().__init__()
 
         self._poller = poller
 
-        self._handlers: ta.Dict[int, FdIoHandler] = {}  # Preserves insertion order
+        self._handlers: ta.Dict[int, FdioHandler] = {}  # Preserves insertion order
 
-    def register(self, h: FdIoHandler) -> None:
+    def register(self, h: FdioHandler) -> None:
         if (hid := id(h)) in self._handlers:
             raise KeyError(h)
         self._handlers[hid] = h
 
-    def unregister(self, h: FdIoHandler) -> None:
+    def unregister(self, h: FdioHandler) -> None:
         del self._handlers[id(h)]
 
     def poll(self, *, timeout: float = 1.) -> None:

omlish/{lite → io}/fdio/pollers.py

@@ -9,7 +9,7 @@ import typing as ta
 ##
 
 
-class FdIoPoller(abc.ABC):
+class FdioPoller(abc.ABC):
     def __init__(self) -> None:
         super().__init__()
 
@@ -120,8 +120,8 @@ class FdIoPoller(abc.ABC):
 ##
 
 
-class SelectFdIoPoller(FdIoPoller):
-    def poll(self, timeout: ta.Optional[float]) -> FdIoPoller.PollResult:
+class SelectFdioPoller(FdioPoller):
+    def poll(self, timeout: ta.Optional[float]) -> FdioPoller.PollResult:
         try:
             r, w, x = select.select(
                 self._readable,
@@ -132,22 +132,22 @@ class SelectFdIoPoller(FdIoPoller):
 
         except OSError as exc:
             if exc.errno == errno.EINTR:
-                return FdIoPoller.PollResult(msg='EINTR encountered in poll', exc=exc)
+                return FdioPoller.PollResult(msg='EINTR encountered in poll', exc=exc)
             elif exc.errno == errno.EBADF:
-                return FdIoPoller.PollResult(msg='EBADF encountered in poll', exc=exc)
+                return FdioPoller.PollResult(msg='EBADF encountered in poll', exc=exc)
             else:
                 raise
 
-        return FdIoPoller.PollResult(r, w)
+        return FdioPoller.PollResult(r, w)
 
 
 ##
 
 
-PollFdIoPoller: ta.Optional[ta.Type[FdIoPoller]]
+PollFdioPoller: ta.Optional[ta.Type[FdioPoller]]
 if hasattr(select, 'poll'):
 
-    class _PollFdIoPoller(FdIoPoller):
+    class _PollFdioPoller(FdioPoller):
         def __init__(self) -> None:
             super().__init__()
 
@@ -180,14 +180,14 @@ if hasattr(select, 'poll'):
 
         #
 
-        def poll(self, timeout: ta.Optional[float]) -> FdIoPoller.PollResult:
+        def poll(self, timeout: ta.Optional[float]) -> FdioPoller.PollResult:
             polled: ta.List[ta.Tuple[int, int]]
             try:
                 polled = self._poller.poll(timeout * 1000 if timeout is not None else None)
 
             except OSError as exc:
                 if exc.errno == errno.EINTR:
-                    return FdIoPoller.PollResult(msg='EINTR encountered in poll', exc=exc)
+                    return FdioPoller.PollResult(msg='EINTR encountered in poll', exc=exc)
                 else:
                     raise
 
@@ -205,8 +205,8 @@ if hasattr(select, 'poll'):
                     r.append(fd)
                 if mask & self._WRITE:
                     w.append(fd)
-            return FdIoPoller.PollResult(r, w, inv=inv)
+            return FdioPoller.PollResult(r, w, inv=inv)
 
-    PollFdIoPoller = _PollFdIoPoller
+    PollFdioPoller = _PollFdioPoller
 else:
-    PollFdIoPoller = None
+    PollFdioPoller = None

omlish/io/generators/__init__.py

@@ -0,0 +1,56 @@
+from .consts import (  # noqa
+    DEFAULT_BUFFER_SIZE,
+)
+
+from .direct import (  # noqa
+    DirectGenerator,
+
+    BytesDirectGenerator,
+    StrDirectGenerator,
+)
+
+from .readers import (  # noqa
+    ReaderGenerator,
+    BytesReaderGenerator,
+    StrReaderGenerator,
+
+    ExactReaderGenerator,
+    BytesExactReaderGenerator,
+    StrExactReaderGenerator,
+
+    GeneratorReader,
+
+    PrependableGeneratorReader,
+    PrependableBytesGeneratorReader,
+    PrependableStrGeneratorReader,
+    prependable_bytes_generator_reader,
+    prependable_str_generator_reader,
+
+    BufferedGeneratorReader,
+    BufferedBytesGeneratorReader,
+    BufferedStrGeneratorReader,
+    buffered_bytes_generator_reader,
+    buffered_str_generator_reader,
+)
+
+from .stepped import (  # noqa
+    SteppedGenerator,
+    BytesSteppedGenerator,
+    StrSteppedGenerator,
+    BytesToStrSteppedGenerator,
+    StrToBytesSteppedGenerator,
+
+    SteppedReaderGenerator,
+    BytesSteppedReaderGenerator,
+    StrSteppedReaderGenerator,
+
+    flatmap_stepped_generator,
+
+    joined_bytes_stepped_generator,
+    joined_str_stepped_generator,
+
+    read_into_bytes_stepped_generator,
+    read_into_str_stepped_generator,
+
+    buffer_bytes_stepped_reader_generator,
+)

omlish/io/generators/consts.py

@@ -0,0 +1 @@
+DEFAULT_BUFFER_SIZE = 4 * 0x1000

omlish/io/generators/direct.py

@@ -0,0 +1,13 @@
+import typing as ta
+
+
+O = ta.TypeVar('O')
+I = ta.TypeVar('I')
+R = ta.TypeVar('R')
+
+
+# Direct generators yield outputs 1:1 with inputs.
+DirectGenerator: ta.TypeAlias = ta.Generator[O, I, R]
+
+BytesDirectGenerator: ta.TypeAlias = DirectGenerator[bytes, bytes, R]
+StrDirectGenerator: ta.TypeAlias = DirectGenerator[str, str, R]

omlish/io/generators/readers.py

@@ -0,0 +1,189 @@
+"""
+TODO:
+ - BufferedBytesGeneratorReader
+ - docstrings
+ - memoryviews
+"""
+import abc
+import typing as ta
+
+from ... import check
+from .consts import DEFAULT_BUFFER_SIZE
+
+
+T = ta.TypeVar('T')
+
+O = ta.TypeVar('O')
+I = ta.TypeVar('I')
+R = ta.TypeVar('R')
+
+AnyT = ta.TypeVar('AnyT', bound=ta.Any)
+
+
+# Reader generators yield ints of amounts of data needed, or None for needing any amount of data.
+ReaderGenerator: ta.TypeAlias = ta.Generator[int | None, I, R]
+
+BytesReaderGenerator: ta.TypeAlias = ReaderGenerator[bytes, R]
+StrReaderGenerator: ta.TypeAlias = ReaderGenerator[str, R]
+
+
+# Exact reader generators always specify read sizes.
+ExactReaderGenerator: ta.TypeAlias = ta.Generator[int, I, R]
+
+BytesExactReaderGenerator: ta.TypeAlias = ExactReaderGenerator[bytes, R]
+StrExactReaderGenerator: ta.TypeAlias = ExactReaderGenerator[str, R]
+
+
+##
+
+
+class _BytesJoiner:
+    def _join(self, lst: list[bytes]) -> bytes:
+        return b''.join(lst)
+
+
+class _StrJoiner:
+    def _join(self, lst: list[str]) -> str:
+        return ''.join(lst)
+
+
+##
+
+
+class GeneratorReader(abc.ABC, ta.Generic[T]):
+    @abc.abstractmethod
+    def read(self, sz: int | None) -> ReaderGenerator[T, T]:
+        raise NotImplementedError
+
+    def read_exact(self, sz: int) -> ReaderGenerator[T, T]:
+        d: ta.Any = yield from self.read(sz)
+        if len(d) != sz:
+            raise EOFError(f'GeneratorReader got {len(d)}, expected {sz}')
+        return d
+
+
+##
+
+
+class PrependableGeneratorReader(GeneratorReader[AnyT]):
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._queue: list[tuple[AnyT, int]] = []
+
+    @abc.abstractmethod
+    def _join(self, lst: list[AnyT]) -> AnyT:
+        raise NotImplementedError
+
+    def read(self, sz: int | None) -> ReaderGenerator[AnyT, AnyT]:
+        if not self._queue:
+            d: AnyT = check.not_none((yield sz))
+            return d
+
+        if sz is None:
+            return self._queue.pop(0)[0]
+
+        lst: list[AnyT] = []
+        rem = sz
+        while rem > 0 and self._queue:
+            c, p = self._queue[0]
+
+            if len(c) - p > rem:
+                lst.append(c[p:p + rem])
+                self._queue[0] = (c, p + rem)
+                return self._join(lst)
+
+            lst.append(c[p:])
+            rem -= len(c) - p
+            self._queue.pop(0)
+
+        if rem:
+            d = check.not_none((yield rem))
+            if d:
+                lst.append(d)  # type: ignore[unreachable]
+
+        if len(lst) == 1:
+            return lst[0]
+        else:
+            return self._join(lst)
+
+    def prepend(self, d: AnyT, p: int | None = None) -> None:
+        if d:
+            self._queue.insert(0, (d, p or 0))
+
+
+class PrependableBytesGeneratorReader(
+    _BytesJoiner,
+    PrependableGeneratorReader[bytes],
+):
+    pass
+
+
+class PrependableStrGeneratorReader(
+    _StrJoiner,
+    PrependableGeneratorReader[str],
+):
+    pass
+
+
+prependable_bytes_generator_reader = PrependableBytesGeneratorReader
+prependable_str_generator_reader = PrependableStrGeneratorReader
+
+
+##
+
+
+class BufferedGeneratorReader(PrependableGeneratorReader[AnyT], abc.ABC):
+    def __init__(
+            self,
+            buffer_size: int = DEFAULT_BUFFER_SIZE,
+    ) -> None:
+        check.arg(buffer_size > 0)
+
+        super().__init__()
+
+        self._buffer_size = buffer_size
+
+    def read(self, sz: int | None) -> ReaderGenerator[AnyT, AnyT]:
+        g = super().read(sz)
+        i: ta.Any = None
+        while True:
+            try:
+                q = g.send(i)
+            except StopIteration as e:
+                return e.value
+
+            check.state(not self._queue)
+
+            if q is None:
+                i = check.not_none((yield None))
+                continue
+
+            r = max(q, self._buffer_size)
+            d: AnyT = check.not_none((yield r))
+            if len(d) < q:
+                i = d
+                continue
+
+            i = d[:q]
+            self.prepend(d, q)
+
+
+class BufferedBytesGeneratorReader(
+    _BytesJoiner,
+    BufferedGeneratorReader[bytes],
+    PrependableGeneratorReader[bytes],
+):
+    pass
+
+
+class BufferedStrGeneratorReader(
+    _StrJoiner,
+    BufferedGeneratorReader[str],
+    PrependableGeneratorReader[str],
+):
+    pass
+
+
+buffered_bytes_generator_reader = BufferedBytesGeneratorReader
+buffered_str_generator_reader = BufferedStrGeneratorReader