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

1from __future__ import annotations 

2 

3from contextvars import ContextVar, Token 

4from dataclasses import dataclass, field 

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

6 

7from .io import FileDescr 

8from .io_basics import FileName 

9from .utils import SLOTS 

10 

11 

12@dataclass(frozen=True, **SLOTS) 

13class PackagingContext: 

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

15 init=False, 

16 default_factory=cast( 

17 "Callable[[], List[Token[Optional[PackagingContext]]]]", list 

18 ), 

19 ) 

20 

21 bioimageio_yaml_file_name: FileName 

22 

23 file_sources: Dict[FileName, FileDescr] 

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

25 

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

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

28 

29 def replace( 

30 self, 

31 *, 

32 bioimageio_yaml_file_name: Optional[FileName] = None, 

33 file_sources: Optional[Dict[FileName, FileDescr]] = None, 

34 weights_priority_order: Union[ 

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

36 ] = "unchanged", 

37 ) -> "PackagingContext": 

38 """return a modiefied copy""" 

39 return PackagingContext( 

40 bioimageio_yaml_file_name=( 

41 self.bioimageio_yaml_file_name 

42 if bioimageio_yaml_file_name is None 

43 else bioimageio_yaml_file_name 

44 ), 

45 file_sources=( 

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

47 ), 

48 weights_priority_order=( 

49 self.weights_priority_order 

50 if weights_priority_order == "unchanged" 

51 else weights_priority_order 

52 ), 

53 ) 

54 

55 def __enter__(self): 

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

57 return self 

58 

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

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

61 

62 

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

64 "packaging_context_var", default=None 

65)