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

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

2 

3import warnings 

4from typing import Any, Callable, List, Optional, Union, cast 

5 

6from pydantic import BaseModel, Field, field_validator, model_validator 

7 

8 

9class PipDeps(BaseModel): 

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

11 

12 pip: List[str] = Field(default_factory=list) 

13 

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

15 @classmethod 

16 def _remove_empty_and_sort(cls, value: List[str]) -> List[str]: 

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

18 

19 def __lt__(self, other: Any): 

20 if isinstance(other, PipDeps): 

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

22 else: 

23 return False 

24 

25 def __gt__(self, other: Any): 

26 if isinstance(other, PipDeps): 

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

28 else: 

29 return False 

30 

31 

32class CondaEnv(BaseModel): 

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

34 

35 name: Optional[str] = None 

36 channels: List[str] = Field(default_factory=list) 

37 dependencies: List[Union[str, PipDeps]] = Field( 

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

39 ) 

40 

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

42 def _ensure_valid_conda_env_name(cls, value: Optional[str]) -> Optional[str]: 

43 if value is None: 

44 return None 

45 

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

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

48 

49 return value or "empty" 

50 

51 @property 

52 def wo_name(self): 

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

54 

55 def _get_version_pin(self, package: str): 

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

57 

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

59 """ 

60 for d in self.dependencies: 

61 if isinstance(d, PipDeps): 

62 for p in d.pip: 

63 if p.startswith(package): 

64 return p[len(package) :] 

65 elif d.startswith(package): 

66 return d[len(package) :] 

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

68 package 

69 ): 

70 return d_wo_channel[len(package) :] 

71 

72 def get_pip_deps(self) -> List[str]: 

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

74 for dep in self.dependencies: 

75 if isinstance(dep, PipDeps): 

76 return dep.pip 

77 

78 return [] 

79 

80 

81class BioimageioCondaEnv(CondaEnv): 

82 """A special `CondaEnv` that 

83 - automatically adds bioimageio specific dependencies 

84 - sorts dependencies 

85 """ 

86 

87 @model_validator(mode="after") 

88 def _normalize_bioimageio_conda_env(self): 

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

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

91 if req_channel not in self.channels: 

92 self.channels.append(req_channel) 

93 

94 if "defaults" in self.channels: 

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

96 self.channels.remove("defaults") 

97 

98 if "pip" not in self.dependencies: 

99 self.dependencies.append("pip") 

100 

101 for dep in self.dependencies: 

102 if isinstance(dep, PipDeps): 

103 pip_section = dep 

104 pip_section.pip.sort() 

105 break 

106 else: 

107 pip_section = None 

108 

109 if ( 

110 pip_section is None 

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

112 ) and not any( 

113 d.startswith("bioimageio.core") 

114 or d.startswith("conda-forge::bioimageio.core") 

115 for d in self.dependencies 

116 if not isinstance(d, PipDeps) 

117 ): 

118 self.dependencies.append("conda-forge::bioimageio.core") 

119 

120 self.dependencies.sort() 

121 return self