omdev 0.0.0.dev221__py3-none-any.whl → 0.0.0.dev223__py3-none-any.whl

omdev/oci/loading.py

@@ -0,0 +1,142 @@
+# ruff: noqa: UP006 UP007
+# @omlish-lite
+import dataclasses as dc
+import json
+import typing as ta
+
+from omlish.lite.check import check
+
+from .data import OciImageConfig
+from .data import OciImageIndex
+from .data import OciImageLayer
+from .data import OciImageManifest
+from .data import is_empty_oci_dataclass
+from .media import OCI_IMAGE_LAYER_KIND_MEDIA_TYPES_
+from .media import OCI_MEDIA_FIELDS
+from .media import OciMediaDescriptor
+from .media import OciMediaImageConfig
+from .media import OciMediaImageIndex
+from .media import OciMediaImageManifest
+from .media import unmarshal_oci_media_dataclass
+from .repositories import FileOciRepository
+from .repositories import OciRepository
+
+
+T = ta.TypeVar('T')
+
+
+##
+
+
+class OciRepositoryLoader:
+    def __init__(
+            self,
+            repo: OciRepository,
+    ) -> None:
+        super().__init__()
+
+        self._repo = repo
+
+    #
+
+    def load_object(
+            self,
+            data: bytes,
+            cls: ta.Type[T] = object,  # type: ignore[assignment]
+            *,
+            media_type: ta.Optional[str] = None,
+    ) -> T:
+        text = data.decode('utf-8')
+        dct = json.loads(text)
+        obj = unmarshal_oci_media_dataclass(
+            dct,
+            media_type=media_type,
+        )
+        return check.isinstance(obj, cls)
+
+    def read_object(
+            self,
+            digest: str,
+            cls: ta.Type[T] = object,  # type: ignore[assignment]
+            *,
+            media_type: ta.Optional[str] = None,
+    ) -> T:
+        data = self._repo.read_blob(digest)
+        return self.load_object(
+            data,
+            cls,
+            media_type=media_type,
+        )
+
+    def read_descriptor(
+            self,
+            desc: OciMediaDescriptor,
+            cls: ta.Type[T] = object,  # type: ignore[assignment]
+    ) -> ta.Any:
+        return self.read_object(
+            desc.digest,
+            cls,
+            media_type=desc.media_type,
+        )
+
+    #
+
+    def from_media(self, obj: ta.Any) -> ta.Any:
+        def make_kw(*exclude):
+            return {
+                a: getattr(obj, a)
+                for f in dc.fields(obj)
+                if (a := f.name) not in OCI_MEDIA_FIELDS
+                and a not in exclude
+            }
+
+        if isinstance(obj, OciMediaImageConfig):
+            return OciImageConfig(**make_kw())
+
+        elif isinstance(obj, OciMediaImageManifest):
+            return OciImageManifest(
+                **make_kw('config', 'layers'),
+                config=self.from_media(self.read_descriptor(obj.config)),
+                layers=[
+                    OciImageLayer(
+                        kind=lk,
+                        data=self._repo.ref_blob(l.digest),
+                    )
+                    for l in obj.layers
+                    if (lk := OCI_IMAGE_LAYER_KIND_MEDIA_TYPES_.get(l.media_type)) is not None
+                ],
+            )
+
+        elif isinstance(obj, OciMediaImageIndex):
+            return OciImageIndex(
+                **make_kw('manifests'),
+                manifests=[
+                    fm
+                    for m in obj.manifests
+                    if self._repo.contains_blob(m.digest)
+                    for fm in [self.from_media(self.read_descriptor(m))]
+                    if not is_empty_oci_dataclass(fm)
+                ],
+            )
+
+        else:
+            raise TypeError(obj)
+
+
+##
+
+
+def read_oci_repository_root_index(
+        obj: ta.Any,
+        *,
+        file_name: str = 'index.json',
+) -> OciImageIndex:
+    file_repo = check.isinstance(OciRepository.of(obj), FileOciRepository)
+
+    repo_ldr = OciRepositoryLoader(file_repo)
+
+    media_image_idx = repo_ldr.load_object(file_repo.read_file(file_name), OciMediaImageIndex)
+
+    image_idx = repo_ldr.from_media(media_image_idx)
+
+    return check.isinstance(image_idx, OciImageIndex)

omdev/oci/media.py

@@ -0,0 +1,179 @@
+# ruff: noqa: UP006 UP007
+# @omlish-lite
+import abc
+import dataclasses as dc
+import typing as ta
+
+from omlish.lite.check import check
+from omlish.lite.marshal import OBJ_MARSHALER_FIELD_KEY
+from omlish.lite.marshal import OBJ_MARSHALER_OMIT_IF_NONE
+from omlish.lite.marshal import unmarshal_obj
+
+from .data import OciImageConfig
+from .data import OciImageLayer
+
+
+##
+
+
+OCI_MEDIA_FIELDS: ta.Collection[str] = frozenset([
+    'schema_version',
+    'media_type',
+])
+
+
+@dc.dataclass()
+class OciMediaDataclass(abc.ABC):  # noqa
+    SCHEMA_VERSION: ta.ClassVar[int]
+
+    @property
+    def schema_version(self) -> int:
+        raise TypeError
+
+    MEDIA_TYPE: ta.ClassVar[str]
+
+    @property
+    def media_type(self) -> str:
+        raise TypeError
+
+    #
+
+    def __init_subclass__(cls, **kwargs: ta.Any) -> None:
+        super().__init_subclass__(**kwargs)
+        for a in OCI_MEDIA_FIELDS:
+            check.in_(a, cls.__dict__)
+
+
+_REGISTERED_OCI_MEDIA_DATACLASSES: ta.Dict[str, ta.Type[OciMediaDataclass]] = {}
+
+
+def _register_oci_media_dataclass(cls):
+    check.issubclass(cls, OciMediaDataclass)
+    check.arg(dc.is_dataclass(cls))
+    mt = check.non_empty_str(cls.__dict__['MEDIA_TYPE'])
+    check.not_in(mt, _REGISTERED_OCI_MEDIA_DATACLASSES)
+    _REGISTERED_OCI_MEDIA_DATACLASSES[mt] = cls
+    return cls
+
+
+def get_registered_oci_media_dataclass(media_type: str) -> ta.Optional[ta.Type[OciMediaDataclass]]:
+    return _REGISTERED_OCI_MEDIA_DATACLASSES.get(media_type)
+
+
+def unmarshal_oci_media_dataclass(
+        dct: ta.Mapping[str, ta.Any],
+        *,
+        media_type: ta.Optional[str] = None,
+) -> ta.Any:
+    if media_type is None:
+        media_type = check.non_empty_str(dct['mediaType'])
+    cls = _REGISTERED_OCI_MEDIA_DATACLASSES[media_type]
+    return unmarshal_obj(dct, cls)
+
+
+##
+
+
+@dc.dataclass()
+class OciMediaDescriptor:
+    """https://github.com/opencontainers/image-spec/blob/92353b0bee778725c617e7d57317b568a7796bd0/descriptor.md#properties"""  # noqa
+
+    media_type: str = dc.field(metadata={OBJ_MARSHALER_FIELD_KEY: 'mediaType'})
+    digest: str
+    size: int
+
+    #
+
+    urls: ta.Optional[ta.Sequence[str]] = dc.field(default=None, metadata={OBJ_MARSHALER_OMIT_IF_NONE: True})
+    annotations: ta.Optional[ta.Mapping[str, str]] = dc.field(default=None, metadata={OBJ_MARSHALER_OMIT_IF_NONE: True})  # noqa
+    data: ta.Optional[str] = dc.field(default=None, metadata={OBJ_MARSHALER_OMIT_IF_NONE: True})
+    artifact_type: ta.Optional[str] = dc.field(default=None, metadata={OBJ_MARSHALER_FIELD_KEY: 'artifactType', OBJ_MARSHALER_OMIT_IF_NONE: True})  # noqa
+
+    #
+
+    platform: ta.Optional[ta.Mapping[str, ta.Any]] = dc.field(default=None, metadata={OBJ_MARSHALER_OMIT_IF_NONE: True})  # noqa
+
+
+##
+
+
+@_register_oci_media_dataclass
+@dc.dataclass()
+class OciMediaImageIndex(OciMediaDataclass):
+    """https://github.com/opencontainers/image-spec/blob/92353b0bee778725c617e7d57317b568a7796bd0/image-index.md"""
+
+    manifests: ta.Sequence[OciMediaDescriptor]  # -> OciMediaImageIndex | OciMediaImageManifest
+
+    #
+
+    annotations: ta.Optional[ta.Mapping[str, str]] = dc.field(default=None, metadata={OBJ_MARSHALER_OMIT_IF_NONE: True})  # noqa
+
+    #
+
+    SCHEMA_VERSION: ta.ClassVar[int] = 2
+    schema_version: int = dc.field(default=SCHEMA_VERSION, metadata={OBJ_MARSHALER_FIELD_KEY: 'schemaVersion'})
+
+    MEDIA_TYPE: ta.ClassVar[str] = 'application/vnd.oci.image.index.v1+json'
+    media_type: str = dc.field(default=MEDIA_TYPE, metadata={OBJ_MARSHALER_FIELD_KEY: 'mediaType'})
+
+
+#
+
+
+@_register_oci_media_dataclass
+@dc.dataclass()
+class OciMediaImageManifest(OciMediaDataclass):
+    """https://github.com/opencontainers/image-spec/blob/92353b0bee778725c617e7d57317b568a7796bd0/manifest.md"""
+
+    config: OciMediaDescriptor  # -> OciMediaImageConfig
+
+    layers: ta.Sequence[OciMediaDescriptor]
+
+    #
+
+    annotations: ta.Optional[ta.Mapping[str, str]] = dc.field(default=None, metadata={OBJ_MARSHALER_OMIT_IF_NONE: True})  # noqa
+
+    #
+
+    SCHEMA_VERSION: ta.ClassVar[int] = 2
+    schema_version: int = dc.field(default=SCHEMA_VERSION, metadata={OBJ_MARSHALER_FIELD_KEY: 'schemaVersion'})
+
+    MEDIA_TYPE: ta.ClassVar[str] = 'application/vnd.oci.image.manifest.v1+json'
+    media_type: str = dc.field(default=MEDIA_TYPE, metadata={OBJ_MARSHALER_FIELD_KEY: 'mediaType'})
+
+
+#
+
+
+OCI_IMAGE_LAYER_KIND_MEDIA_TYPES: ta.Mapping[OciImageLayer.Kind, str] = {
+    OciImageLayer.Kind.TAR: 'application/vnd.oci.image.layer.v1.tar',
+    OciImageLayer.Kind.TAR_GZIP: 'application/vnd.oci.image.layer.v1.tar+gzip',
+    OciImageLayer.Kind.TAR_ZSTD: 'application/vnd.oci.image.layer.v1.tar+zstd',
+}
+
+OCI_IMAGE_LAYER_KIND_MEDIA_TYPES_: ta.Mapping[str, OciImageLayer.Kind] = {
+    v: k
+    for k, v in OCI_IMAGE_LAYER_KIND_MEDIA_TYPES.items()
+}
+
+
+#
+
+
+@_register_oci_media_dataclass
+@dc.dataclass()
+class OciMediaImageConfig(OciImageConfig, OciMediaDataclass):
+    SCHEMA_VERSION: ta.ClassVar[int] = 2
+    schema_version: int = dc.field(default=SCHEMA_VERSION, metadata={OBJ_MARSHALER_FIELD_KEY: 'schemaVersion'})
+
+    MEDIA_TYPE: ta.ClassVar[str] = 'application/vnd.oci.image.config.v1+json'
+    media_type: str = dc.field(default=MEDIA_TYPE, metadata={OBJ_MARSHALER_FIELD_KEY: 'mediaType'})
+
+
+##
+
+
+OCI_MANIFEST_MEDIA_TYPES: ta.AbstractSet[str] = frozenset([
+    OciMediaImageIndex.MEDIA_TYPE,
+    OciMediaImageManifest.MEDIA_TYPE,
+])

omdev/oci/packing.py

@@ -0,0 +1,381 @@
+# ruff: noqa: UP006 UP007
+# @omlish-lite
+import contextlib
+import heapq
+import os.path
+import tarfile
+import typing as ta
+
+from omlish.lite.cached import cached_nullary
+from omlish.lite.check import check
+from omlish.lite.contextmanagers import ExitStacked
+
+from .compression import OciCompression
+from .tars import OciDataTarWriter
+from .tars import WrittenOciDataTarFileInfo
+
+
+##
+
+
+class OciLayerUnpacker(ExitStacked):
+    def __init__(
+            self,
+            input_files: ta.Sequence[ta.Union[str, tarfile.TarFile]],
+            output_file_path: str,
+    ) -> None:
+        super().__init__()
+
+        self._input_files = list(input_files)
+        self._output_file_path = output_file_path
+
+    #
+
+    @contextlib.contextmanager
+    def _open_input_file(self, input_file: ta.Union[str, tarfile.TarFile]) -> ta.Iterator[tarfile.TarFile]:
+        if isinstance(input_file, tarfile.TarFile):
+            yield input_file
+
+        elif isinstance(input_file, str):
+            with tarfile.open(input_file) as tar_file:
+                yield tar_file
+
+        else:
+            raise TypeError(input_file)
+
+    #
+
+    class _Entry(ta.NamedTuple):
+        file: ta.Union[str, tarfile.TarFile]
+        info: tarfile.TarInfo
+
+    def _build_input_file_sorted_entries(self, input_file: ta.Union[str, tarfile.TarFile]) -> ta.Sequence[_Entry]:
+        dct: ta.Dict[str, OciLayerUnpacker._Entry] = {}
+
+        with self._open_input_file(input_file) as input_tar_file:
+            for info in input_tar_file.getmembers():
+                check.not_in(info.name, dct)
+                dct[info.name] = self._Entry(
+                    file=input_file,
+                    info=info,
+                )
+
+        return sorted(dct.values(), key=lambda entry: entry.info.name)
+
+    @cached_nullary
+    def _entries_by_name(self) -> ta.Mapping[str, _Entry]:
+        root: dict = {}
+
+        def find_dir(dir_name: str) -> dict:  # noqa
+            if dir_name:
+                dir_parts = dir_name.split('/')
+            else:
+                dir_parts = []
+
+            cur = root  # noqa
+            for dir_part in dir_parts:
+                cur = cur[dir_part]  # noqa
+
+            return check.isinstance(cur, dict)
+
+        #
+
+        for input_file in self._input_files:
+            sorted_entries = self._build_input_file_sorted_entries(input_file)
+
+            wh_names = set()
+            wh_opaques = set()
+
+            #
+
+            for entry in sorted_entries:
+                info = entry.info
+                name = check.non_empty_str(info.name)
+                base_name = os.path.basename(name)
+                dir_name = os.path.dirname(name)
+
+                if base_name == '.wh..wh..opq':
+                    wh_opaques.add(dir_name)
+                    continue
+
+                if base_name.startswith('.wh.'):
+                    wh_base_name = os.path.basename(base_name[4:])
+                    wh_name = os.path.join(dir_name, wh_base_name)
+                    wh_names.add(wh_name)
+                    continue
+
+                cur = find_dir(dir_name)
+
+                if info.type == tarfile.DIRTYPE:
+                    try:
+                        ex = cur[base_name]
+                    except KeyError:
+                        cur[base_name] = {'': entry}
+                    else:
+                        ex[''] = entry
+
+                else:
+                    cur[base_name] = entry
+
+            #
+
+            for wh_name in reversed(sorted(wh_names)):  # noqa
+                wh_dir_name = os.path.dirname(wh_name)
+                wh_base_name = os.path.basename(wh_name)
+
+                cur = find_dir(wh_dir_name)
+                rm = cur[wh_base_name]
+
+                if isinstance(rm, dict):
+                    # Whiteouts wipe out whole directory:
+                    # https://github.com/containerd/containerd/blob/59c8cf6ea5f4175ad512914dd5ce554942bf144f/pkg/archive/tar_test.go#L648
+                    # check.equal(set(rm), '')
+                    del cur[wh_base_name]
+
+                elif isinstance(rm, self._Entry):
+                    del cur[wh_base_name]
+
+                else:
+                    raise TypeError(rm)
+
+            if wh_opaques:
+                raise NotImplementedError
+
+        #
+
+        out: ta.Dict[str, OciLayerUnpacker._Entry] = {}
+
+        def rec(cur):  # noqa
+            for _, child in sorted(cur.items(), key=lambda t: t[0]):
+                if isinstance(child, dict):
+                    rec(child)
+
+                elif isinstance(child, self._Entry):
+                    check.not_in(child.info.name, out)
+                    out[child.info.name] = child
+
+                else:
+                    raise TypeError(child)
+
+        rec(root)
+
+        return out
+
+    #
+
+    @cached_nullary
+    def _output_tar_file(self) -> tarfile.TarFile:
+        return self._enter_context(tarfile.open(self._output_file_path, 'w'))
+
+    #
+
+    def _add_unpacked_entry(
+            self,
+            input_tar_file: tarfile.TarFile,
+            info: tarfile.TarInfo,
+    ) -> None:
+        base_name = os.path.basename(info.name)
+        check.state(not base_name.startswith('.wh.'))
+
+        if info.type in tarfile.REGULAR_TYPES:
+            with check.not_none(input_tar_file.extractfile(info)) as f:
+                self._output_tar_file().addfile(info, f)
+
+        else:
+            self._output_tar_file().addfile(info)
+
+    def _unpack_file(
+            self,
+            input_file: ta.Union[str, tarfile.TarFile],
+    ) -> None:
+        entries_by_name = self._entries_by_name()
+
+        with self._open_input_file(input_file) as input_tar_file:
+            info: tarfile.TarInfo
+            for info in input_tar_file.getmembers():
+                try:
+                    entry = entries_by_name[info.name]
+                except KeyError:
+                    continue
+
+                if entry.file != input_file:
+                    continue
+
+                self._add_unpacked_entry(input_tar_file, info)
+
+    @cached_nullary
+    def write(self) -> None:
+        for input_file in self._input_files:
+            self._unpack_file(input_file)
+
+
+#
+
+
+class OciLayerPacker(ExitStacked):
+    def __init__(
+            self,
+            input_file_path: str,
+            output_file_paths: ta.Sequence[str],
+            *,
+            compression: ta.Optional[OciCompression] = None,
+    ) -> None:
+        super().__init__()
+
+        self._input_file_path = input_file_path
+        self._output_file_paths = list(output_file_paths)
+        self._compression = compression
+
+        self._output_file_indexes_by_name: ta.Dict[str, int] = {}
+
+    #
+
+    @cached_nullary
+    def _input_tar_file(self) -> tarfile.TarFile:
+        # FIXME: check uncompressed
+        return self._enter_context(tarfile.open(self._input_file_path))
+
+    #
+
+    @cached_nullary
+    def _entries_by_name(self) -> ta.Mapping[str, tarfile.TarInfo]:
+        return {
+            info.name: info
+            for info in self._input_tar_file().getmembers()
+        }
+
+    #
+
+    class _CategorizedEntries(ta.NamedTuple):
+        files_by_name: ta.Mapping[str, tarfile.TarInfo]
+        non_files_by_name: ta.Mapping[str, tarfile.TarInfo]
+        links_by_name: ta.Mapping[str, tarfile.TarInfo]
+
+    @cached_nullary
+    def _categorized_entries(self) -> _CategorizedEntries:
+        files_by_name: ta.Dict[str, tarfile.TarInfo] = {}
+        non_files_by_name: ta.Dict[str, tarfile.TarInfo] = {}
+        links_by_name: ta.Dict[str, tarfile.TarInfo] = {}
+
+        for name, info in self._entries_by_name().items():
+            if info.type in tarfile.REGULAR_TYPES:
+                files_by_name[name] = info
+            elif info.type in (tarfile.LNKTYPE, tarfile.GNUTYPE_LONGLINK):
+                links_by_name[name] = info
+            else:
+                non_files_by_name[name] = info
+
+        return self._CategorizedEntries(
+            files_by_name=files_by_name,
+            non_files_by_name=non_files_by_name,
+            links_by_name=links_by_name,
+        )
+
+    #
+
+    @cached_nullary
+    def _non_files_sorted_by_name(self) -> ta.Sequence[tarfile.TarInfo]:
+        return sorted(
+            self._categorized_entries().non_files_by_name.values(),
+            key=lambda info: info.name,
+        )
+
+    @cached_nullary
+    def _files_descending_by_size(self) -> ta.Sequence[tarfile.TarInfo]:
+        return sorted(
+            self._categorized_entries().files_by_name.values(),
+            key=lambda info: -check.isinstance(info.size, int),
+        )
+
+    #
+
+    @cached_nullary
+    def _output_files(self) -> ta.Sequence[ta.BinaryIO]:
+        return [
+            self._enter_context(open(output_file_path, 'wb'))
+            for output_file_path in self._output_file_paths
+        ]
+
+    @cached_nullary
+    def _output_tar_writers(self) -> ta.Sequence[OciDataTarWriter]:
+        return [
+            self._enter_context(
+                OciDataTarWriter(
+                    output_file,
+                    compression=self._compression,
+                ),
+            )
+            for output_file in self._output_files()
+        ]
+
+    #
+
+    def _write_entry(
+            self,
+            info: tarfile.TarInfo,
+            output_file_idx: int,
+    ) -> None:
+        check.not_in(info.name, self._output_file_indexes_by_name)
+
+        writer = self._output_tar_writers()[output_file_idx]
+
+        if info.type in tarfile.REGULAR_TYPES:
+            with check.not_none(self._input_tar_file().extractfile(info)) as f:
+                writer.add_file(info, f)  # type: ignore
+
+        else:
+            writer.add_file(info)
+
+        self._output_file_indexes_by_name[info.name] = output_file_idx
+
+    @cached_nullary
+    def _write_non_files(self) -> None:
+        for non_file in self._non_files_sorted_by_name():
+            self._write_entry(non_file, 0)
+
+    @cached_nullary
+    def _write_files(self) -> None:
+        writers = self._output_tar_writers()
+
+        bins = [
+            (writer.info().compressed_sz, i)
+            for i, writer in enumerate(writers)
+        ]
+
+        heapq.heapify(bins)
+
+        for file in self._files_descending_by_size():
+            _, bin_index = heapq.heappop(bins)
+
+            writer = writers[bin_index]
+
+            self._write_entry(file, bin_index)
+
+            bin_size = writer.info().compressed_sz
+
+            heapq.heappush(bins, (bin_size, bin_index))
+
+    @cached_nullary
+    def _write_links(self) -> None:
+        for link in self._categorized_entries().links_by_name.values():
+            link_name = check.non_empty_str(link.linkname)
+
+            output_file_idx = self._output_file_indexes_by_name[link_name]
+
+            self._write_entry(link, output_file_idx)
+
+    @cached_nullary
+    def write(self) -> ta.Mapping[str, WrittenOciDataTarFileInfo]:
+        writers = self._output_tar_writers()
+
+        self._write_non_files()
+        self._write_files()
+        self._write_links()
+
+        for output_tar_writer in writers:
+            output_tar_writer.tar_file().close()
+
+        return {
+            output_file_path: output_tar_writer.info()
+            for output_file_path, output_tar_writer in zip(self._output_file_paths, writers)
+        }