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

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