Coverage for install/scipp/core/structured.py: 60%

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

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

3# @author Simon Heybrock 

4from collections.abc import Iterator 

5from typing import Any 

6 

7from .._scipp.core import _element_keys, _get_elements, _set_elements 

8from .cpp_classes import DataArray, Dataset, DType, DTypeError, Variable 

9 

10 

11def _prop(key: str) -> property: 

12 def getter(self: Any) -> Variable: 

13 return _get_elements(self._var, key) # type: ignore[no-any-return] 

14 

15 def setter(self: Any, x: Variable) -> None: 

16 _set_elements(self._var, key, x) 

17 

18 return property(getter, setter) 

19 

20 

21def is_structured(obj: Variable | DataArray) -> bool: 

22 """Check whether a variable has a structured dtype.""" 

23 if obj.bins is not None: 

24 data = obj.bins.constituents['data'] 

25 if isinstance(data, Dataset): 

26 raise DTypeError( 

27 "Datasets cannot have a structured dtype, " 

28 "got a variable with Datasets in bins." 

29 ) 

30 return is_structured(data) 

31 return obj.dtype in [DType.vector3, DType.linear_transform3] 

32 

33 

34class Fields: 

35 """Accessor for fields of a structured variable.""" 

36 

37 def __init__(self, keys: list[str], var: Variable) -> None: 

38 self._var = var 

39 self._keys = keys 

40 

41 def __contains__(self, key: str) -> bool: 

42 return key in self._keys 

43 

44 def __iter__(self) -> Iterator[str]: 

45 return self._keys.__iter__() 

46 

47 def __getitem__(self, key: str) -> Variable: 

48 return _get_elements(self._var, key) # type: ignore[no-any-return] 

49 

50 def __setitem__(self, key: str, x: Variable) -> None: 

51 _set_elements(self._var, key, x) 

52 

53 def keys(self) -> Iterator[str]: 

54 yield from self._keys 

55 

56 def values(self) -> Iterator[Variable]: 

57 yield from (self[key] for key in self) 

58 

59 def items(self) -> Iterator[tuple[str, Variable]]: 

60 yield from ((key, self[key]) for key in self) 

61 

62 

63# Make a new 'Fields' type to dynamically insert properties. 

64# This means that the return annotation has to be `Any` and not `Fields`, 

65# because otherwise, using the properties would result in type errors. 

66def _fields(obj: Variable) -> Any: 

67 fields = type('Fields', Fields.__bases__, dict(Fields.__dict__)) 

68 if is_structured(obj): 

69 keys = _element_keys(obj) 

70 for key in keys: 

71 setattr(fields, key, _prop(key)) 

72 return fields(keys=keys, var=obj) 

73 return None