Coverage for src/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 local_files_only: bool = False 

30 """Whether to include only local files when packaging. If True, remote files will be excluded.""" 

31 

32 def replace( 

33 self, 

34 *, 

35 bioimageio_yaml_file_name: Optional[FileName] = None, 

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

37 weights_priority_order: Union[ 

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

39 ] = "unchanged", 

40 local_files_only: Optional[bool] = None, 

41 ) -> "PackagingContext": 

42 """return a modiefied copy""" 

43 return PackagingContext( 

44 bioimageio_yaml_file_name=( 

45 self.bioimageio_yaml_file_name 

46 if bioimageio_yaml_file_name is None 

47 else bioimageio_yaml_file_name 

48 ), 

49 file_sources=( 

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

51 ), 

52 weights_priority_order=( 

53 self.weights_priority_order 

54 if weights_priority_order == "unchanged" 

55 else weights_priority_order 

56 ), 

57 local_files_only=( 

58 self.local_files_only if local_files_only is None else local_files_only 

59 ), 

60 ) 

61 

62 def __enter__(self): 

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

64 return self 

65 

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

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

68 

69 

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

71 "packaging_context_var", default=None 

72)