bioimageio.spec.conda_env

  1import warnings
  2from typing import Any, Callable, List, Optional, Union, cast
  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(
 36        default_factory=cast(Callable[[], List[Union[str, PipDeps]]], list)
 37    )
 38
 39    @field_validator("name", mode="after")
 40    def _ensure_valid_conda_env_name(cls, value: Optional[str]) -> Optional[str]:
 41        if value is None:
 42            return None
 43
 44        for illegal in ("/", " ", ":", "#"):
 45            value = value.replace(illegal, "")
 46
 47        return value or "empty"
 48
 49    @property
 50    def wo_name(self):
 51        return self.model_construct(**{k: v for k, v in self if k != "name"})
 52
 53    def _get_version(self, package: str):
 54        """Helper to return any verison pin for **package**
 55
 56        TODO: improve: interprete version pin and return structured information.
 57        """
 58        for d in self.dependencies:
 59            if isinstance(d, PipDeps):
 60                for p in d.pip:
 61                    if p.startswith(package):
 62                        return p[len(package) :]
 63            elif d.startswith(package):
 64                return d[len(package) :]
 65
 66    def get_pip_deps(self) -> List[str]:
 67        """Get the pip dependencies of this conda env."""
 68        for dep in self.dependencies:
 69            if isinstance(dep, PipDeps):
 70                return dep.pip
 71
 72        return []
 73
 74
 75class BioimageioCondaEnv(CondaEnv):
 76    """A special `CondaEnv` that
 77    - automatically adds bioimageio specific dependencies
 78    - sorts dependencies
 79    """
 80
 81    @model_validator(mode="after")
 82    def _normalize_bioimageio_conda_env(self):
 83        """update a conda env such that we have bioimageio.core and sorted dependencies"""
 84        for req_channel in ("conda-forge", "nodefaults"):
 85            if req_channel not in self.channels:
 86                self.channels.append(req_channel)
 87
 88        if "defaults" in self.channels:
 89            warnings.warn("removing 'defaults' from conda-channels")
 90            self.channels.remove("defaults")
 91
 92        if "pip" not in self.dependencies:
 93            self.dependencies.append("pip")
 94
 95        for dep in self.dependencies:
 96            if isinstance(dep, PipDeps):
 97                pip_section = dep
 98                pip_section.pip.sort()
 99                break
100        else:
101            pip_section = None
102
103        if (
104            pip_section is None
105            or not any(pd.startswith("bioimageio.core") for pd in pip_section.pip)
106        ) and not any(
107            d.startswith("bioimageio.core")
108            or d.startswith("conda-forge::bioimageio.core")
109            for d in self.dependencies
110            if not isinstance(d, PipDeps)
111        ):
112            self.dependencies.append("conda-forge::bioimageio.core")
113
114        self.dependencies.sort()
115        return self