Coverage for src/bioimageio/spec/conda_env.py: 68%

1"""Representation of conda environment.yaml files for bioimageio specifications.""" 

2 

3from __future__ import annotations 

4 

5import warnings 

6from typing import Any, Callable, List, Union, cast 

7 

8from pydantic import BaseModel, Field, field_validator, model_validator 

9 

10 

11class PipDeps(BaseModel): 

12 """Pip dependencies to include in conda dependecies""" 

13 

14 pip: list[str] = Field(default_factory=list) 

15 

16 @field_validator("pip", mode="after") 

17 @classmethod 

18 def _remove_empty_and_sort(cls, value: list[str]) -> list[str]: 

19 return sorted(vs for v in value if (vs := v.strip())) 

20 

21 def __lt__(self, other: Any): 

22 if isinstance(other, PipDeps): 

23 return len(self.pip) < len(other.pip) 

24 else: 

25 return False 

26 

27 def __gt__(self, other: Any): 

28 if isinstance(other, PipDeps): 

29 return len(self.pip) > len(other.pip) 

30 else: 

31 return False 

32 

33 

34class CondaEnv(BaseModel): 

35 """Represenation of the content of a conda environment.yaml file""" 

36 

37 name: str | None = None 

38 channels: list[str] = Field(default_factory=list) 

39 dependencies: list[str | PipDeps] = Field( 

40 default_factory=cast(Callable[[], List[Union[str, PipDeps]]], list) 

41 ) 

42 

43 @field_validator("name", mode="after") 

44 def _ensure_valid_conda_env_name(cls, value: str | None) -> str | None: 

45 if value is None: 

46 return None 

47 

48 for illegal in ("/", " ", ":", "#"): 

49 value = value.replace(illegal, "") 

50 

51 return value or "empty" 

52 

53 @property 

54 def wo_name(self): 

55 return self.model_construct(**{k: v for k, v in self if k != "name"}) 

56 

57 def _get_version_pin(self, package: str): 

58 """Helper to return any version pin for **package** 

59 

60 TODO: improve: interprete version pin and return structured information. 

61 """ 

62 for d in self.dependencies: 

63 if isinstance(d, PipDeps): 

64 for p in d.pip: 

65 if p.startswith(package): 

66 return p[len(package) :] 

67 elif d.startswith(package): 

68 return d[len(package) :] 

69 elif "::" in d and (d_wo_channel := d.split("::", 1)[-1]).startswith( 

70 package 

71 ): 

72 return d_wo_channel[len(package) :] 

73 

74 def get_pip_deps(self) -> list[str]: 

75 """Get the pip dependencies of this conda env.""" 

76 for dep in self.dependencies: 

77 if isinstance(dep, PipDeps): 

78 return dep.pip 

79 

80 return [] 

81 

82 

83class BioimageioCondaEnv(CondaEnv): 

84 """A special `CondaEnv` that 

85 - automatically adds bioimageio specific dependencies 

86 - sorts dependencies 

87 """ 

88 

89 @model_validator(mode="after") 

90 def _normalize_bioimageio_conda_env(self): 

91 """update a conda env such that we have bioimageio.core and sorted dependencies""" 

92 for req_channel in ("conda-forge", "nodefaults"): 

93 if req_channel not in self.channels: 

94 self.channels.append(req_channel) 

95 

96 if "defaults" in self.channels: 

97 warnings.warn("removing 'defaults' from conda-channels") 

98 self.channels.remove("defaults") 

99 

100 if "pip" not in self.dependencies: 

101 self.dependencies.append("pip") 

102 

103 for dep in self.dependencies: 

104 if isinstance(dep, PipDeps): 

105 pip_section = dep 

106 pip_section.pip.sort() 

107 break 

108 else: 

109 pip_section = None 

110 

111 if ( 

112 pip_section is None 

113 or not any(pd.startswith("bioimageio.core") for pd in pip_section.pip) 

114 ) and not any( 

115 d.startswith(("bioimageio.core", "conda-forge::bioimageio.core")) 

116 for d in self.dependencies 

117 if not isinstance(d, PipDeps) 

118 ): 

119 self.dependencies.append("conda-forge::bioimageio.core>=0.9.4") 

120 

121 self.dependencies.sort() 

122 return self