improve extract_resource_package

and use spec.resolve_rdf_source_and_type

Author: FynnBe

bioimageio/core/resource_io/common.py

@@ -1,48 +1,98 @@
 import os
 import pathlib
+import shutil
 import warnings
+import zipfile
 from copy import deepcopy
-from typing import Any, Dict, Optional, Sequence, Tuple, Union
+from typing import Dict, IO, Optional, Sequence, Tuple, Union
 from zipfile import ZIP_DEFLATED, ZipFile
 
-from marshmallow import ValidationError, missing
+from marshmallow import missing
 
 from bioimageio import spec
 from bioimageio.core.resource_io.nodes import ResourceDescription
+from bioimageio.spec.io_ import resolve_rdf_source
 from bioimageio.spec.shared import raw_nodes
-from bioimageio.spec.shared.common import get_class_name_from_type
+from bioimageio.spec.shared.common import BIOIMAGEIO_CACHE_PATH, get_class_name_from_type
 from bioimageio.spec.shared.raw_nodes import ResourceDescription as RawResourceDescription
 from bioimageio.spec.shared.utils import PathToRemoteUriTransformer
 from . import nodes
-from .common import BIOIMAGEIO_CACHE_PATH, yaml
-from .utils import _download_uri_to_local_path, resolve_local_uri, resolve_raw_resource_description, resolve_uri
+from .utils import resolve_raw_resource_description, resolve_uri
 
-
-ROOT_PATH = "root_path"
 serialize_raw_resource_description = spec.io_.serialize_raw_resource_description
 save_raw_resource_description = spec.io_.save_raw_resource_description
 
 
+def extract_resource_package(
+    source: Union[os.PathLike, IO, str, bytes, raw_nodes.URI]
+) -> Tuple[dict, str, pathlib.Path]:
+    """extract a zip source to BIOIMAGEIO_CACHE_PATH"""
+    source, source_name, root = resolve_rdf_source(source)
+    if isinstance(root, bytes):
+        raise NotImplementedError("package source was bytes")
+
+    cache_folder = BIOIMAGEIO_CACHE_PATH / "extracted_packages"
+    cache_folder.mkdir(exist_ok=True, parents=True)
+
+    if isinstance(root, raw_nodes.URI):
+        from urllib.request import urlretrieve
+
+        package_path = cache_folder / root.scheme / root.authority / root.path.strip("/") / root.query
+        if (package_path / "rdf.yaml").exists():
+            download = None
+        else:
+            download, header = urlretrieve(str(root))
+
+        local_source = download
+    else:
+        download = None
+        local_source = root
+        package_path = cache_folder / root.relative_to(list(root.parents)[-1])
+
+    if local_source is not None:
+        with zipfile.ZipFile(local_source) as zf:
+            zf.extractall(package_path)
+
+    if not (package_path / "rdf.yaml").exists():
+        raise FileNotFoundError(f"missing 'rdf.yaml' in {root} extracted from {download}")
+
+    if download is not None:
+        try:
+            os.remove(download)
+        except Exception as e:
+            warnings.warn(f"Could not remove download {download} due to {e}")
+
+    assert isinstance(package_path, pathlib.Path)
+    return source, source_name, package_path
+
+
 def _replace_relative_paths_for_remote_source(
-    raw_rd: RawResourceDescription, source: Union[Any, str, raw_nodes.URI]
+    raw_rd: RawResourceDescription, root: Union[pathlib.Path, raw_nodes.URI, bytes]
 ) -> RawResourceDescription:
-    if isinstance(source, raw_nodes.URI) or isinstance(source, str) and source.startswith("http"):
+    if isinstance(root, raw_nodes.URI):
         # for a remote source relative paths are invalid; replace all relative file paths in source with URLs
-        if isinstance(source, str):
-            source = raw_nodes.URI(source)
-
         warnings.warn(
-            f"changing file paths in RDF to URIs due to a remote {source.scheme} source "
+            f"changing file paths in RDF to URIs due to a remote {root.scheme} source "
             "(may result in an invalid node)"
         )
-        raw_rd = PathToRemoteUriTransformer(remote_source=source).transform(raw_rd)
-        raw_rd.root_path = pathlib.Path()  # root_path cannot be URI
+        raw_rd = PathToRemoteUriTransformer(remote_source=root).transform(raw_rd)
+        root_path = pathlib.Path()  # root_path cannot be URI
+    elif isinstance(root, pathlib.Path):
+        if zipfile.is_zipfile(root):
+            _, _, root_path = extract_resource_package(root)
+        else:
+            root_path = root
+    elif isinstance(root, bytes):
+        raise NotImplementedError("root as bytes (io)")
+    else:
+        raise TypeError(root)
 
+    raw_rd.root_path = root_path
     return raw_rd
 
 
 def load_raw_resource_description(
-    source: Union[os.PathLike, str, dict, raw_nodes.URI, RawResourceDescription]
+    source: Union[dict, os.PathLike, IO, str, bytes, raw_nodes.URI]
 ) -> RawResourceDescription:
     """load a raw python representation from a BioImage.IO resource description file (RDF).
     Use `load_resource_description` for a more convenient representation.
@@ -53,12 +103,8 @@ def load_raw_resource_description(
     Returns:
         raw BioImage.IO resource
     """
-    if isinstance(source, RawResourceDescription):
-        return source
-
-    data, type_ = resolve_rdf_source_and_type(source)
-    raw_rd = spec.load_raw_resource_description(data, update_to_current_format=True)
-    raw_rd = _replace_relative_paths_for_remote_source(raw_rd, source)
+    raw_rd = spec.load_raw_resource_description(source, update_to_current_format=True)
+    raw_rd = _replace_relative_paths_for_remote_source(raw_rd, raw_rd.root_path)
     return raw_rd
 
 
@@ -186,26 +232,6 @@ def _get_tmp_package_path(raw_rd: RawResourceDescription, weights_priority_order
     return package_path
 
 
-def extract_resource_package(source: Union[os.PathLike, str, raw_nodes.URI]) -> pathlib.Path:
-    """extract a zip source to BIOIMAGEIO_CACHE_PATH"""
-    local_source = resolve_uri(source)
-    assert isinstance(local_source, pathlib.Path)
-    cache_folder = BIOIMAGEIO_CACHE_PATH / "extracted_packages"
-    cache_folder.mkdir(exist_ok=True, parents=True)
-    package_path = cache_folder / f"{local_source.stem}"
-    with ZipFile(local_source) as zf:
-        zf.extractall(package_path)
-
-    for rdf_name in ["rdf.yaml", "model.yaml", "rdf.yml", "model.yml"]:
-        rdf_path = package_path / rdf_name
-        if rdf_path.exists():
-            break
-    else:
-        raise FileNotFoundError(local_source / "rdf.yaml")
-
-    return rdf_path
-
-
 def make_zip(
     path: os.PathLike, content: Dict[str, Union[str, pathlib.Path]], *, compression: int, compression_level: int
 ) -> None:
@@ -225,53 +251,3 @@ def make_zip(
                 myzip.writestr(arc_name, file_or_str_content)
             else:
                 myzip.write(file_or_str_content, arcname=arc_name)
-
-
-def resolve_rdf_source_and_type(source: Union[os.PathLike, str, dict, raw_nodes.URI]) -> Tuple[dict, str]:
-    if isinstance(source, dict):
-        data = source
-        if ROOT_PATH not in data:
-            data[ROOT_PATH] = pathlib.Path()
-    else:
-        data = get_dict_from_yaml_source(source)
-
-    type_ = data.get("type", "model")  # todo: remove default 'model' type
-
-    return data, type_
-
-
-def get_dict_from_yaml_source(source: Union[os.PathLike, str, raw_nodes.URI, dict]) -> dict:
-    if isinstance(source, dict):
-        if ROOT_PATH not in source:
-            source[ROOT_PATH] = pathlib.Path()
-
-        return source
-    elif isinstance(source, (str, os.PathLike, raw_nodes.URI)):
-        source = resolve_local_uri(source, pathlib.Path())
-    else:
-        raise TypeError(source)
-
-    if isinstance(source, raw_nodes.URI):  # remote uri
-        local_source = _download_uri_to_local_path(source)
-        root_path = pathlib.Path()
-    else:
-        local_source = source
-        root_path = source.parent
-
-    assert isinstance(local_source, pathlib.Path)
-    if local_source.suffix == ".zip":
-        local_source = extract_resource_package(local_source)
-        root_path = local_source.parent
-
-    if local_source.suffix == ".yml":
-        warnings.warn(
-            "suffix '.yml' is not recommended and will raise a ValidationError in the future. Use '.yaml' instead "
-            "(https://yaml.org/faq.html)"
-        )
-    elif local_source.suffix != ".yaml":
-        raise ValidationError(f"invalid suffix {local_source.suffix} for source {source}")
-
-    data = yaml.load(local_source)
-    assert isinstance(data, dict)
-    data[ROOT_PATH] = root_path
-    return data

bioimageio/core/resource_io/io_.py

@@ -14,9 +14,9 @@
 from marshmallow import ValidationError
 
 from bioimageio.spec.shared import fields, raw_nodes
+from bioimageio.spec.shared.common import BIOIMAGEIO_CACHE_PATH
 from bioimageio.spec.shared.utils import GenericRawNode, GenericRawRD, NodeTransformer, NodeVisitor
 from . import nodes
-from .common import BIOIMAGEIO_CACHE_PATH
 
 GenericResolvedNode = typing.TypeVar("GenericResolvedNode", bound=nodes.Node)
 GenericNode = typing.Union[GenericRawNode, GenericResolvedNode]

bioimageio/core/resource_io/utils.py

@@ -6,6 +6,11 @@
 from bioimageio.core import load_resource_description
 
 
+def test_validate_model(unet2d_nuclei_broad_model):
+    ret = subprocess.run(["bioimageio", "validate", unet2d_nuclei_broad_model])
+    assert ret.returncode == 0
+
+
 def test_cli_test_model(unet2d_nuclei_broad_model):
     ret = subprocess.run(["bioimageio", "test-model", unet2d_nuclei_broad_model])
     assert ret.returncode == 0