Coverage for bioimageio/spec/_internal/packaging_context.py: 95%

1from __future__ import annotations 

2 

3from contextvars import ContextVar, Token 

4from dataclasses import dataclass, field 

5from typing import Dict, List, Literal, Optional, Sequence, Union 

6 

7from .io_basics import AbsoluteFilePath, FileName, ZipPath 

8from .url import HttpUrl 

9 

10 

11@dataclass(frozen=True) 

12class PackagingContext: 

13 _context_tokens: "List[Token[Optional[PackagingContext]]]" = field( 

14 init=False, default_factory=list 

15 ) 

16 

17 bioimageio_yaml_file_name: FileName 

18 

19 file_sources: Dict[FileName, Union[AbsoluteFilePath, HttpUrl, ZipPath]] 

20 """File sources to include in the packaged resource""" 

21 

22 weights_priority_order: Optional[Sequence[str]] = None 

23 """set to select a single weights entry when packaging model resources""" 

24 

25 def replace( 

26 self, 

27 *, 

28 bioimageio_yaml_file_name: Optional[FileName] = None, 

29 file_sources: Optional[ 

30 Dict[FileName, Union[AbsoluteFilePath, HttpUrl, ZipPath]] 

31 ] = None, 

32 weights_priority_order: Union[ 

33 Optional[Sequence[str]], Literal["unchanged"] 

34 ] = "unchanged", 

35 ) -> "PackagingContext": 

36 """return a modiefied copy""" 

37 return PackagingContext( 

38 bioimageio_yaml_file_name=( 

39 self.bioimageio_yaml_file_name 

40 if bioimageio_yaml_file_name is None 

41 else bioimageio_yaml_file_name 

42 ), 

43 file_sources=( 

44 dict(self.file_sources) if file_sources is None else file_sources 

45 ), 

46 weights_priority_order=( 

47 self.weights_priority_order 

48 if weights_priority_order == "unchanged" 

49 else weights_priority_order 

50 ), 

51 ) 

52 

53 def __enter__(self): 

54 self._context_tokens.append(packaging_context_var.set(self)) 

55 return self 

56 

57 def __exit__(self, type, value, traceback): # type: ignore 

58 packaging_context_var.reset(self._context_tokens.pop(-1)) 

59 

60 

61packaging_context_var: ContextVar[Optional[PackagingContext]] = ContextVar( 

62 "packaging_context_var", default=None 

63)