Coverage for arrakis/channel.py: 93.9%

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 expected_latency: int, optional

44 Expected publication latency for this channel's data, in

45 seconds.

46

47 """

48

49 name: str

50 data_type: numpy.dtype | str

51 sample_rate: float

52 time: int | None = None

53 publisher: str | None = None

54 partition_id: str | None = None

55 expected_latency: int | None = None

56

57 @property

58 def dtype(self):

59 return self.data_type

60

61 def __post_init__(self) -> None:

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

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

64 self.validate()

65

66 def validate(self) -> None:

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

68 if len(components) != 2:

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

70 raise ValueError(msg)

71

72 def __repr__(self) -> str:

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

74

75 def __str__(self) -> str:

76 return self.name

77

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

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

80 is_equal = (

81 self.name == other.name

82 and self.dtype == other.dtype

83 and self.sample_rate == other.sample_rate

84 )

85

86 # optional fields match only if both are defined

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

88 is_equal &= self.time == other.time

89 if self.publisher and other.publisher:

90 is_equal &= self.publisher == other.publisher

91 if self.partition_id and other.partition_id:

92 is_equal &= self.partition_id == other.partition_id

93

94 return is_equal

95

96 @cached_property

97 def domain(self) -> str:

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

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

100

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

102 """Serialize channel metadata to JSON.

103

104 Parameters

105 ----------

106 time : int, optional

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

108

109 """

110 # generate dict from dataclass and adjust fields

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

112 # channel name, as well as updating the timestamp

113 # if passed in.

114 obj = asdict(self)

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

116 if time is not None:

117 obj["time"] = time

118 obj = {k: v for k, v in obj.items() if v is not None}

119 return json.dumps(obj)

120

121 @classmethod

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

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

124

125 Parameters

126 ----------

127 payload : str

128 The JSON-serialized channel.

129

130 Returns

131 -------

132 Channel

133 The newly created channel.

134

135 """

136 obj = json.loads(payload)

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

138 return cls(**obj)

139

140 @classmethod

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

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

143

144 Parameters

145 ----------

146 field : pyarrow.field

147 The channel field containing relevant metadata.

148

149 Returns

150 -------

151 Channel

152 The newly created channel.

153

154 """

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

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

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

158

159

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

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

162

163 Note that this does not always match the string representation

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

165 different languages.

166

167 Parameters

168 ----------

169 dtype : pyarrow.ListType

170 The list data type to inspect.

171

172 Returns

173 -------

174 str

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

176

177 """

178 inner_dtype = str(dtype.value_type)

179 if inner_dtype == "float":

180 return "float32"

181 if inner_dtype == "double":

182 return "float64"

183 return inner_dtype