Coverage for src / bioimageio / core / _collection.py: 0%

1"""Collection utilities for bioimageio""" 

2 

3import json 

4import warnings 

5from functools import cache 

6from typing import Literal, Mapping, Sequence, TypedDict 

7 

8from bioimageio.spec.common import Sha256 

9from bioimageio.spec.utils import get_reader 

10 

11from ._settings import settings 

12 

13 

14class IndexItemVersion(TypedDict): 

15 version: str 

16 source: str 

17 sha256: Sha256 

18 

19 

20class IndexItem(TypedDict): 

21 id: str 

22 type: str 

23 versions: Sequence[IndexItemVersion] 

24 

25 

26class Index(TypedDict): 

27 items: Sequence[IndexItem] 

28 total: int 

29 count_per_type: Mapping[str, int] 

30 timestamp: str 

31 

32 

33class IdPartsEntry(TypedDict): 

34 nouns: Mapping[str, str] 

35 adjectives: Sequence[str] 

36 

37 

38class CollectionConfig(TypedDict): 

39 id_parts: Mapping[Literal["model", "dataset", "notebook"], IdPartsEntry] 

40 reviewers: Mapping[str, Sequence[str]] 

41 

42 

43@cache 

44def load_json(url: str): 

45 reader = get_reader(url) 

46 return json.load(reader) 

47 

48 

49def load_index() -> Index: 

50 return load_json(settings.collection_index_url) 

51 

52 

53def load_collection_config() -> CollectionConfig: 

54 return load_json(settings.collection_config_url) 

55 

56 

57@cache 

58def load_hypened_nouns() -> set[str]: 

59 """get all nouns with hyphens that could be part of a nickname, e.g. 't-rex'""" 

60 return { 

61 noun 

62 for id_parts in load_collection_config()["id_parts"].values() 

63 for noun in id_parts.get("nouns", []) 

64 if "-" in noun 

65 } 

66 

67 

68def lookup_from_index(source: str) -> tuple[str, dict[Literal["sha256"], Sha256]]: 

69 index = load_index() 

70 if source.startswith("bioimage-io/") and source.count("/") == 2: 

71 version = source.split("/")[-1] 

72 source = source[: -(len(version) + 1)] 

73 else: 

74 version = None 

75 

76 for item in index["items"]: 

77 if item["id"] in (source, f"bioimage-io/{source}"): 

78 v = item["versions"][-1] 

79 if version is not None: 

80 for v in item["versions"]: 

81 if v["version"] == version: 

82 break 

83 else: 

84 warnings.warn( 

85 f"Version {version} not found in index, using latest version." 

86 ) 

87 else: 

88 continue 

89 

90 return v["source"], {"sha256": v["sha256"]} 

91 

92 return source, {} 

93 

94 

95def get_resource_icon(nickname: str, rtype: str) -> str: 

96 """Get emoji for a resource, matching to its nickname noun. nicknames are of the form "{adjective}-{noun}", e.g. "affable-shark".""" 

97 if "-" not in nickname: 

98 return " " 

99 

100 # remove hyphen from noun part of nickname, e.g. "laid-back-t-rex" -> "laid-back-trex" 

101 for hyphened_noun in load_hypened_nouns(): 

102 if nickname.endswith(hyphened_noun): 

103 nickname = nickname[: -len(hyphened_noun)] + hyphened_noun.replace("-", "") 

104 

105 # last hyphen now sparates adjective and noun, e.g. "laid-back-trex" -> "laid-back" and "trex" 

106 noun = nickname[nickname.rfind("-") + 1 :].replace("-", "") 

107 try: 

108 ret = { 

109 k.replace("-", ""): v 

110 for k, v in load_collection_config()["id_parts"][ 

111 rtype if rtype in ("model", "dataset", "notebook") else "notebook" 

112 ] 

113 .get("nouns", {}) 

114 .items() 

115 }.get(noun, " ") 

116 except Exception as e: 

117 warnings.warn(f"Error getting icon for {rtype} {nickname}: {e}") 

118 ret = " " 

119 

120 return ret