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

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 Dict, Literal, Optional, Union 

7 

8from ..typing import DataArray, Dataset, Variable 

9from .argument_handlers import combine_dict_args 

10 

11 

12def _assign( 

13 obj: Union[Dataset, DataArray], 

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

15 obj_attrs: Optional[Dict[str, Variable]] = None, 

16 /, 

17 **kw_obj_attrs, 

18) -> Union[Dataset, DataArray]: 

19 out = obj.copy(deep=False) 

20 collected = combine_dict_args(obj_attrs, kw_obj_attrs) 

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

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

23 return out 

24 

25 

26def assign_coords( 

27 self, coords: Optional[Dict[str, Variable]] = None, /, **coords_kwargs 

28) -> Union[DataArray, Dataset]: 

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

30 

31 Parameters 

32 ---------- 

33 coords : 

34 New coordinates. 

35 

36 coords_kwargs : 

37 Keyword arguments form of ``coords``. 

38 

39 Returns 

40 ------- 

41 : 

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

43 

44 """ 

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

46 

47 

48def assign_masks( 

49 self, masks: Optional[Dict[str, Variable]] = None, /, **masks_kwargs 

50) -> DataArray: 

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

52 

53 Parameters 

54 ---------- 

55 masks : 

56 New masks. 

57 

58 masks_kwargs : 

59 Keyword arguments form of ``masks``. 

60 

61 Returns 

62 ------- 

63 : 

64 ``scipp.DataArray`` with updated masks. 

65 

66 """ 

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

68 

69 

70def assign_attrs( 

71 self, attrs: Optional[Dict[str, Variable]] = None, /, **attrs_kwargs 

72) -> DataArray: 

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

74 

75 Parameters 

76 ---------- 

77 attrs : 

78 New attrs. 

79 

80 attrs_kwargs : 

81 Keyword arguments form of ``attrs``. 

82 

83 Returns 

84 ------- 

85 : 

86 ``scipp.DataArray`` with updated attributes. 

87 

88 """ 

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