bioimageio.spec.conda_env

  1import warnings
  2from typing import Any, List, Optional, Union
  3
  4from pydantic import BaseModel, Field, field_validator, model_validator
  5
  6
  7class PipDeps(BaseModel):
  8    """Pip dependencies to include in conda dependecies"""
  9
 10    pip: List[str] = Field(default_factory=list)
 11
 12    def __lt__(self, other: Any):
 13        if isinstance(other, PipDeps):
 14            return len(self.pip) < len(other.pip)
 15        else:
 16            return False
 17
 18    def __gt__(self, other: Any):
 19        if isinstance(other, PipDeps):
 20            return len(self.pip) > len(other.pip)
 21        else:
 22            return False
 23
 24
 25class CondaEnv(BaseModel):
 26    """Represenation of the content of a conda environment.yaml file"""
 27
 28    name: Optional[str] = None
 29    channels: List[str] = Field(default_factory=list)
 30    dependencies: List[Union[str, PipDeps]] = Field(default_factory=list)
 31
 32    @field_validator("name", mode="after")
 33    def _ensure_valid_conda_env_name(cls, value: Optional[str]) -> Optional[str]:
 34        if value is None:
 35            return None
 36
 37        for illegal in ("/", " ", ":", "#"):
 38            value = value.replace(illegal, "")
 39
 40        return value or "empty"
 41
 42    @property
 43    def wo_name(self):
 44        return self.model_construct(**{k: v for k, v in self if k != "name"})
 45
 46    def _get_version(self, package: str):
 47        """Helper to return any verison pin for **package**
 48
 49        TODO: improve: interprete version pin and return structured information.
 50        """
 51        for d in self.dependencies:
 52            if isinstance(d, PipDeps):
 53                for p in d.pip:
 54                    if p.startswith(package):
 55                        return p[len(package) :]
 56            elif d.startswith(package):
 57                return d[len(package) :]
 58
 59
 60class BioimageioCondaEnv(CondaEnv):
 61    """A special `CondaEnv` that
 62    - automatically adds bioimageio specific dependencies
 63    - sorts dependencies
 64    """
 65
 66    @model_validator(mode="after")
 67    def _normalize_bioimageio_conda_env(self):
 68        """update a conda env such that we have bioimageio.core and sorted dependencies"""
 69        for req_channel in ("conda-forge", "nodefaults"):
 70            if req_channel not in self.channels:
 71                self.channels.append(req_channel)
 72
 73        if "defaults" in self.channels:
 74            warnings.warn("removing 'defaults' from conda-channels")
 75            self.channels.remove("defaults")
 76
 77        if "pip" not in self.dependencies:
 78            self.dependencies.append("pip")
 79
 80        for dep in self.dependencies:
 81            if isinstance(dep, PipDeps):
 82                pip_section = dep
 83                pip_section.pip.sort()
 84                break
 85        else:
 86            pip_section = None
 87
 88        if (
 89            pip_section is not None
 90            and any(pd.startswith("bioimageio.core") for pd in pip_section.pip)
 91        ) and not any(
 92            d.startswith("bioimageio.core")
 93            or d.startswith("conda-forge::bioimageio.core")
 94            for d in self.dependencies
 95            if not isinstance(d, PipDeps)
 96        ):
 97            self.dependencies.append("conda-forge::bioimageio.core")
 98
 99        self.dependencies.sort()
100        return self