Coverage for install/scipp/core/assignments.py: 47%

1# SPDX-License-Identifier: BSD-3-Clause 

2# Copyright (c) 2023 Scipp contributors (https://github.com/scipp) 

3 

4from __future__ import annotations 

5 

6from typing import Literal, TypeVar 

7 

8from .argument_handlers import combine_dict_args 

9from .cpp_classes import DataArray, Dataset, Variable 

10 

11_T = TypeVar('_T', Dataset, DataArray) 

12 

13 

14def _assign( 

15 obj: _T, 

16 name: Literal['coords', 'masks', 'attrs'], 

17 obj_attrs: dict[str, Variable] | None = None, 

18 /, 

19 **kw_obj_attrs: Variable, 

20) -> _T: 

21 out = obj.copy(deep=False) 

22 collected = combine_dict_args(obj_attrs, kw_obj_attrs) 

23 for key, value in collected.items(): 

24 getattr(out, name)[key] = value 

25 return out 

26 

27 

28def assign_coords( 

29 self: _T, coords: dict[str, Variable] | None = None, /, **coords_kwargs: Variable 

30) -> _T: 

31 """Return new object with updated or inserted coordinate. 

32 

33 Parameters 

34 ---------- 

35 coords : 

36 New coordinates. 

37 

38 coords_kwargs : 

39 Keyword arguments form of ``coords``. 

40 

41 Returns 

42 ------- 

43 : 

44 ``scipp.DataArray`` or ``scipp.Dataset`` with updated coordinates. 

45 

46 """ 

47 return _assign(self, 'coords', coords, **coords_kwargs) 

48 

49 

50def assign_masks( 

51 self: DataArray, 

52 masks: dict[str, Variable] | None = None, 

53 /, 

54 **masks_kwargs: Variable, 

55) -> DataArray: 

56 """Return new object with updated or inserted masks. 

57 

58 Parameters 

59 ---------- 

60 masks : 

61 New masks. 

62 

63 masks_kwargs : 

64 Keyword arguments form of ``masks``. 

65 

66 Returns 

67 ------- 

68 : 

69 ``scipp.DataArray`` with updated masks. 

70 

71 """ 

72 return _assign(self, 'masks', masks, **masks_kwargs) 

73 

74 

75def assign_attrs( 

76 self: DataArray, 

77 attrs: dict[str, Variable] | None = None, 

78 /, 

79 **attrs_kwargs: Variable, 

80) -> DataArray: 

81 """Return new object with updated or inserted attrs. 

82 

83 Parameters 

84 ---------- 

85 attrs : 

86 New attrs. 

87 

88 attrs_kwargs : 

89 Keyword arguments form of ``attrs``. 

90 

91 Returns 

92 ------- 

93 : 

94 ``scipp.DataArray`` with updated attributes. 

95 

96 """ 

97 return _assign(self, 'attrs', attrs, **attrs_kwargs)