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_pin(self, package: str):
 54        """Helper to return any version 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            elif "::" in d and (d_wo_channel := d.split("::", 1)[-1]).startswith(
 66                package
 67            ):
 68                return d_wo_channel[len(package) :]
 69
 70    def get_pip_deps(self) -> List[str]:
 71        """Get the pip dependencies of this conda env."""
 72        for dep in self.dependencies:
 73            if isinstance(dep, PipDeps):
 74                return dep.pip
 75
 76        return []
 77
 78
 79class BioimageioCondaEnv(CondaEnv):
 80    """A special `CondaEnv` that
 81    - automatically adds bioimageio specific dependencies
 82    - sorts dependencies
 83    """
 84
 85    @model_validator(mode="after")
 86    def _normalize_bioimageio_conda_env(self):
 87        """update a conda env such that we have bioimageio.core and sorted dependencies"""
 88        for req_channel in ("conda-forge", "nodefaults"):
 89            if req_channel not in self.channels:
 90                self.channels.append(req_channel)
 91
 92        if "defaults" in self.channels:
 93            warnings.warn("removing 'defaults' from conda-channels")
 94            self.channels.remove("defaults")
 95
 96        if "pip" not in self.dependencies:
 97            self.dependencies.append("pip")
 98
 99        for dep in self.dependencies:
100            if isinstance(dep, PipDeps):
101                pip_section = dep
102                pip_section.pip.sort()
103                break
104        else:
105            pip_section = None
106
107        if (
108            pip_section is None
109            or not any(pd.startswith("bioimageio.core") for pd in pip_section.pip)
110        ) and not any(
111            d.startswith("bioimageio.core")
112            or d.startswith("conda-forge::bioimageio.core")
113            for d in self.dependencies
114            if not isinstance(d, PipDeps)
115        ):
116            self.dependencies.append("conda-forge::bioimageio.core")
117
118        self.dependencies.sort()
119        return self