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