introduce modelpack for model in OCI format

oras/defaults.py

@@ -48,3 +48,6 @@ class registry:
 blank_config_hash = (
     "sha256:44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a"
 )
+
+# ModelPack annotations, refer to https://github.com/modelpack/model-spec/blob/main/specs-go/v1/annotations.go
+annotation_filepath = "org.cnai.model.filepath"

oras/provider.py

@@ -802,6 +802,12 @@ def push(
             if annotations:
                 layer["annotations"].update(annotations)
 
+            # if filepath annotation is present, set it
+            if oras.defaults.annotation_filepath in layer["annotations"]:
+                layer["annotations"][oras.defaults.annotation_filepath] = (
+                    blob_name.strip(os.sep)
+                )
+
             # update the manifest with the new layer
             manifest["layers"].append(layer)
             logger.debug(f"Preparing layer {layer}")
@@ -901,8 +907,8 @@ def pull(
         files = []
         for layer in manifest.get("layers", []):
             filename = (layer.get("annotations") or {}).get(
-                oras.defaults.annotation_title
-            )
+                oras.defaults.annotation_filepath
+            ) or (layer.get("annotations") or {}).get(oras.defaults.annotation_title)
 
             # If we don't have a filename, default to digest. Hopefully does not happen
             if not filename:
@@ -917,18 +923,32 @@ def pull(
                 )
                 continue
 
-            # A directory will need to be uncompressed and moved
-            if layer["mediaType"] == oras.defaults.default_blob_dir_media_type:
-                targz = oras.utils.get_tmpfile(suffix=".tar.gz")
-                self.download_blob(container, layer["digest"], targz)
-
-                # The artifact will be extracted to the correct name
-                oras.utils.extract_targz(targz, os.path.dirname(outfile))
+            # Determine compression format from mediaType and handle accordingly
+            media_type = layer["mediaType"]
+            compression = oras.utils.get_compression_from_media_type(media_type)
+
+            # Handle compressed archives that need extraction
+            if compression in ["gzip", "zstd", "tar"]:
+                # Get appropriate file extension based on compression
+                if compression == "gzip":
+                    suffix = ".tar.gz"
+                elif compression == "zstd":
+                    suffix = ".tar.zst"
+                else:  # tar
+                    suffix = ".tar"
+
+                archive_file = oras.utils.get_tmpfile(suffix=suffix)
+                self.download_blob(container, layer["digest"], archive_file)
+
+                # Extract the archive to the correct location
+                oras.utils.extract_by_compression(
+                    archive_file, os.path.dirname(outfile), compression
+                )
 
             # Anything else just extracted directly
             else:
                 self.download_blob(container, layer["digest"], outfile)
-            logger.info(f"Successfully pulled {outfile}.")
+            logger.info(f"Successfully pulled {outfile}")
             files.append(outfile)
         return files
 

oras/utils/__init__.py

@@ -1,6 +1,10 @@
 from .fileio import (
     copyfile,
+    extract_by_compression,
+    extract_tar,
+    extract_tar_zstd,
     extract_targz,
+    get_compression_from_media_type,
     get_file_hash,
     get_size,
     get_tmpdir,

oras/utils/fileio.py

@@ -103,6 +103,92 @@ def extract_targz(targz: str, outdir: str, numeric_owner: bool = False):
         tar.extractall(outdir, members=None, numeric_owner=numeric_owner)
 
 
+def extract_tar(tar_file: str, outdir: str, numeric_owner: bool = False):
+    """
+    Extract a .tar (uncompressed) to an output directory.
+    """
+    with tarfile.open(tar_file, "r:") as tar:
+        for member in tar.getmembers():
+            member_path = os.path.join(outdir, member.name)
+            if not is_within_directory(outdir, member_path):
+                raise Exception("Attempted Path Traversal in Tar File")
+        tar.extractall(outdir, members=None, numeric_owner=numeric_owner)
+
+
+def extract_tar_zstd(tar_zstd: str, outdir: str, numeric_owner: bool = False):
+    """
+    Extract a .tar.zst/.tar.zstd to an output directory.
+    Requires zstandard package to be installed.
+    """
+    try:
+        import zstandard as zstd
+    except ImportError:
+        raise ImportError(
+            "zstandard package is required for zstd decompression. "
+            "Install it with: pip install zstandard"
+        )
+
+    with open(tar_zstd, "rb") as compressed_file:
+        dctx = zstd.ZstdDecompressor()
+        with dctx.stream_reader(compressed_file) as reader:
+            with tarfile.open(fileobj=reader, mode="r|") as tar:
+                for member in tar:
+                    member_path = os.path.join(outdir, member.name)
+                    if not is_within_directory(outdir, member_path):
+                        raise Exception("Attempted Path Traversal in Tar File")
+                    tar.extract(member, outdir, numeric_owner=numeric_owner)
+
+
+def get_compression_from_media_type(media_type: str) -> str:
+    """
+    Determine compression format from media type suffix.
+
+    :param media_type: The media type string
+    :type media_type: str
+    :return: Compression format ('gzip', 'zstd', 'tar', or 'raw')
+    :rtype: str
+    """
+    if media_type.endswith("+gzip"):
+        return "gzip"
+    elif media_type.endswith("+zstd"):
+        return "zstd"
+    elif media_type.endswith(".tar"):
+        return "tar"
+    elif media_type.endswith(".raw"):
+        return "raw"
+    else:
+        # Default to gzip for backward compatibility
+        return "gzip"
+
+
+def extract_by_compression(
+    archive_file: str, outdir: str, compression: str, numeric_owner: bool = False
+):
+    """
+    Extract archive based on compression format.
+
+    :param archive_file: Path to the archive file
+    :type archive_file: str
+    :param outdir: Output directory for extraction
+    :type outdir: str
+    :param compression: Compression format ('gzip', 'zstd', 'tar', or 'raw')
+    :type compression: str
+    :param numeric_owner: Whether to use numeric owner
+    :type numeric_owner: bool
+    """
+    if compression == "gzip":
+        extract_targz(archive_file, outdir, numeric_owner)
+    elif compression == "zstd":
+        extract_tar_zstd(archive_file, outdir, numeric_owner)
+    elif compression == "tar":
+        extract_tar(archive_file, outdir, numeric_owner)
+    elif compression == "raw":
+        # For raw files, no extraction needed - they should be handled differently
+        raise ValueError("Raw files should not be extracted as archives")
+    else:
+        raise ValueError(f"Unsupported compression format: {compression}")
+
+
 def is_within_directory(directory: str, target: str) -> bool:
     """
     Determine whether a file is within a directory