Coverage for arrakis/channel.py: 90.6%

1# Copyright (c) 2022, California Institute of Technology and contributors

2#

3# You should have received a copy of the licensing terms for this

4# software included in the file "LICENSE" located in the top-level

5# directory of this package. If you did not, you can view a copy at

6# https://git.ligo.org/ngdd/arrakis-python/-/raw/main/LICENSE

7

8"""Channel information."""

9

10 from __future__ import annotations

11

12 import json

13 from dataclasses import asdict, dataclass

14 from functools import cached_property

15 from typing import TYPE_CHECKING

16

17 import numpy

18

19 if TYPE_CHECKING:

20 import pyarrow

21

22

23@dataclass(frozen=True)

24 class Channel:

25 """Metadata associated with a channel.

26

27 Channels have the form {domain}:*.

28

29 Parameters

30 ----------

31 name : str

32 The name associated with this channel.

33 data_type : numpy.dtype

34 The data type associated with this channel.

35 sample_rate : float

36 The sampling rate associated with this channel.

37 time : int, optional

38 The timestamp when this metadata became active.

39 publisher : str

40 The publisher associated with this channel.

41 partition_id : str, optional

42 The partition ID associated with this channel.

43

44 """

45

46 name: str

47 data_type: numpy.dtype

48 sample_rate: float

49 time: int | None = None

50 publisher: str | None = None

51 partition_id: str | None = None

52

53 @property

54 def dtype(self):

55 return self.data_type

56

57 def __post_init__(self) -> None:

58 # cast to numpy dtype object, as raw types like numpy.float64 are not

59 self.validate()

60 object.__setattr__(self, "data_type", numpy.dtype(self.data_type))

61

62 def validate(self) -> None:

63 components = self.name.split(":")

64 if len(components) != 2:

65 msg = "channel is malformed, needs to be in the form {domain}:*"

66 raise ValueError(msg)

67

68 def __repr__(self) -> str:

69 return f"<{self.name}, {self.sample_rate} Hz, {self.data_type}>"

70

71 def __str__(self) -> str:

72 return self.name

73

74 def __eq__(self, other) -> bool:

75 # name, data type and sample rate are always required to match

76 is_equal = (

77 self.name == other.name

78 and self.dtype == other.dtype

79 and self.sample_rate == other.sample_rate

80 )

81

82 # optional fields match only if both are defined

83 if self.time is not None and other.time is not None:

84 is_equal &= self.time == other.time

85 if self.publisher and other.publisher:

86 is_equal &= self.publisher == other.publisher

87 if self.partition_id and other.partition_id:

88 is_equal &= self.partition_id == other.partition_id

89

90 return is_equal

91

92 @cached_property

93 def domain(self) -> str:

94 """The domain associated with this channel."""

95 return self.name.split(":", 1)[0]

96

97 def to_json(self, time: int | None = None) -> str:

98 """Serialize channel metadata to JSON.

99

100 Parameters

101 ----------

102 time : int, optional

103 If specified, the timestamp when this metadata became active.

104

105 """

106 # generate dict from dataclass and adjust fields

107 # to be JSON compatible. In addition, store the

108 # channel name, as well as updating the timestamp

109 # if passed in.

110 obj = asdict(self)

111 obj["data_type"] = numpy.dtype(self.data_type).name

112 if time is not None:

113 obj["time"] = time

114 return json.dumps(obj)

115

116 @classmethod

117 def from_json(cls, payload: str) -> Channel:

118 """Create a Channel from its JSON representation.

119

120 Parameters

121 ----------

122 payload : str

123 The JSON-serialized channel.

124

125 Returns

126 -------

127 Channel

128 The newly created channel.

129

130 """

131 obj = json.loads(payload)

132 obj["data_type"] = numpy.dtype(obj["data_type"])

133 return cls(**obj)

134

135 @classmethod

136 def from_field(cls, field: pyarrow.field) -> Channel:

137 """Create a Channel from Arrow Flight field metadata.

138

139 Parameters

140 ----------

141 field : pyarrow.field

142 The channel field containing relevant metadata.

143

144 Returns

145 -------

146 Channel

147 The newly created channel.

148

149 """

150 data_type = numpy.dtype(_list_dtype_to_str(field.type))

151 sample_rate = float(field.metadata[b"rate"].decode())

152 return cls(field.name, data_type, sample_rate)

153

154

155 def _list_dtype_to_str(dtype: pyarrow.ListType) -> str:

156 """Return a string representation of the list's inner data type.

157

158 Note that this does not always match the string representation

159 of Arrow's internal data types, to match the behavior across

160 different languages.

161

162 Parameters

163 ----------

164 dtype : pyarrow.ListType

165 The list data type to inspect.

166

167 Returns

168 -------

169 str

170 A string representation of the list's inner data type.

171

172 """

173 inner_dtype = str(dtype.value_type)

174 if inner_dtype == "float":

175 return "float32"

176 if inner_dtype == "double":

177 return "float64"

178 return inner_dtype