Explicitly mention `black` in CodingStyle guidelines
[babeltrace.git] / src / bindings / python / bt2 / bt2 / trace_collection_message_iterator.py
CommitLineData
d34e69cf
PP
1# The MIT License (MIT)
2#
3# Copyright (c) 2017 Philippe Proulx <pproulx@efficios.com>
4#
5# Permission is hereby granted, free of charge, to any person obtaining a copy
6# of this software and associated documentation files (the "Software"), to deal
7# in the Software without restriction, including without limitation the rights
8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9# copies of the Software, and to permit persons to whom the Software is
10# furnished to do so, subject to the following conditions:
11#
12# The above copyright notice and this permission notice shall be included in
13# all copies or substantial portions of the Software.
14#
15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
21# THE SOFTWARE.
22
23from bt2 import utils
24import bt2
88fdcc33 25import itertools
fa4c33e3 26import bt2.message_iterator
d34e69cf 27import datetime
d34e69cf
PP
28from collections import namedtuple
29import numbers
30
31
88fdcc33
PP
32# a pair of component and ComponentSpec
33_ComponentAndSpec = namedtuple('_ComponentAndSpec', ['comp', 'spec'])
d34e69cf
PP
34
35
88fdcc33 36class ComponentSpec:
cc81b5ab
PP
37 def __init__(self, plugin_name, class_name, params=None,
38 logging_level=bt2.logging.LoggingLevel.NONE):
d34e69cf 39 utils._check_str(plugin_name)
c88be1c8 40 utils._check_str(class_name)
cc81b5ab 41 utils._check_log_level(logging_level)
d34e69cf 42 self._plugin_name = plugin_name
c88be1c8 43 self._class_name = class_name
cc81b5ab 44 self._logging_level = logging_level
d34e69cf
PP
45
46 if type(params) is str:
da35796c 47 self._params = bt2.create_value({'paths': [params]})
d34e69cf
PP
48 else:
49 self._params = bt2.create_value(params)
50
51 @property
52 def plugin_name(self):
53 return self._plugin_name
54
55 @property
c88be1c8
PP
56 def class_name(self):
57 return self._class_name
d34e69cf 58
cc81b5ab
PP
59 @property
60 def logging_level(self):
61 return self._logging_level
62
d34e69cf
PP
63 @property
64 def params(self):
65 return self._params
66
67
68# datetime.datetime or integral to nanoseconds
69def _get_ns(obj):
70 if obj is None:
71 return
72
73 if isinstance(obj, numbers.Real):
74 # consider that it's already in seconds
75 s = obj
76 elif isinstance(obj, datetime.datetime):
77 # s -> ns
78 s = obj.timestamp()
79 else:
80 raise TypeError('"{}" is not an integral number or a datetime.datetime object'.format(obj))
81
82 return int(s * 1e9)
83
84
88fdcc33
PP
85class _CompClsType:
86 SOURCE = 0
87 FILTER = 1
88
89
fa4c33e3 90class TraceCollectionMessageIterator(bt2.message_iterator._MessageIterator):
88fdcc33 91 def __init__(self, source_component_specs, filter_component_specs=None,
da35796c 92 stream_intersection_mode=False, begin=None, end=None):
d34e69cf
PP
93 utils._check_bool(stream_intersection_mode)
94 self._stream_intersection_mode = stream_intersection_mode
95 self._begin_ns = _get_ns(begin)
96 self._end_ns = _get_ns(end)
88fdcc33
PP
97
98 if type(source_component_specs) is ComponentSpec:
99 source_component_specs = [source_component_specs]
100
101 if type(filter_component_specs) is ComponentSpec:
102 filter_component_specs = [filter_component_specs]
103 elif filter_component_specs is None:
104 filter_component_specs = []
105
d34e69cf 106 self._src_comp_specs = source_component_specs
88fdcc33 107 self._flt_comp_specs = filter_component_specs
d34e69cf
PP
108 self._next_suffix = 1
109 self._connect_ports = False
110
88fdcc33 111 # lists of _ComponentAndSpec
d34e69cf 112 self._src_comps_and_specs = []
88fdcc33 113 self._flt_comps_and_specs = []
d34e69cf 114
88fdcc33
PP
115 self._validate_component_specs(source_component_specs)
116 self._validate_component_specs(filter_component_specs)
d34e69cf
PP
117 self._build_graph()
118
88fdcc33
PP
119 def _validate_component_specs(self, comp_specs):
120 for comp_spec in comp_specs:
121 if type(comp_spec) is not ComponentSpec:
122 raise TypeError('"{}" object is not a ComponentSpec'.format(type(comp_spec)))
d34e69cf
PP
123
124 def __next__(self):
fa4c33e3 125 return next(self._msg_iter)
d34e69cf 126
da35796c 127 def _create_stream_intersection_trimmer(self, component, port):
d34e69cf
PP
128 # find the original parameters specified by the user to create
129 # this port's component to get the `path` parameter
130 for src_comp_and_spec in self._src_comps_and_specs:
da35796c 131 if component == src_comp_and_spec.comp:
d34e69cf
PP
132 break
133
134 try:
da35796c
SM
135 paths = src_comp_and_spec.spec.params['paths']
136 except Exception as e:
137 raise bt2.Error('all source components must be created with a "paths" parameter in stream intersection mode') from e
d34e69cf 138
da35796c 139 params = {'paths': paths}
d34e69cf
PP
140
141 # query the port's component for the `trace-info` object which
142 # contains the stream intersection range for each exposed
143 # trace
144 query_exec = bt2.QueryExecutor()
c88be1c8 145 trace_info_res = query_exec.query(src_comp_and_spec.comp.cls,
d34e69cf
PP
146 'trace-info', params)
147 begin = None
148 end = None
149
ddf49b27 150 # find the trace info for this port's trace
f4811b4f
PP
151 try:
152 for trace_info in trace_info_res:
ddf49b27
SM
153 for stream in trace_info['streams']:
154 if stream['port-name'] == port.name:
155 range_ns = trace_info['intersection-range-ns']
156 begin = range_ns['begin']
157 end = range_ns['end']
158 break
da35796c 159 except Exception:
f4811b4f 160 pass
d34e69cf
PP
161
162 if begin is None or end is None:
163 raise bt2.Error('cannot find stream intersection range for port "{}"'.format(port.name))
164
da35796c 165 name = 'trimmer-{}-{}'.format(src_comp_and_spec.comp.name, port.name)
d34e69cf
PP
166 return self._create_trimmer(begin, end, name)
167
168 def _create_muxer(self):
169 plugin = bt2.find_plugin('utils')
170
171 if plugin is None:
172 raise bt2.Error('cannot find "utils" plugin (needed for the muxer)')
173
174 if 'muxer' not in plugin.filter_component_classes:
175 raise bt2.Error('cannot find "muxer" filter component class in "utils" plugin')
176
177 comp_cls = plugin.filter_component_classes['muxer']
178 return self._graph.add_component(comp_cls, 'muxer')
179
da35796c 180 def _create_trimmer(self, begin_ns, end_ns, name):
d34e69cf
PP
181 plugin = bt2.find_plugin('utils')
182
183 if plugin is None:
184 raise bt2.Error('cannot find "utils" plugin (needed for the trimmer)')
185
186 if 'trimmer' not in plugin.filter_component_classes:
187 raise bt2.Error('cannot find "trimmer" filter component class in "utils" plugin')
188
189 params = {}
190
da35796c
SM
191 def ns_to_string(ns):
192 s_part = ns // 1000000000
193 ns_part = ns % 1000000000
194 return '{}.{:09d}'.format(s_part, ns_part)
d34e69cf 195
da35796c
SM
196 if begin_ns is not None:
197 params['begin'] = ns_to_string(begin_ns)
198
199 if end_ns is not None:
200 params['end'] = ns_to_string(end_ns)
d34e69cf
PP
201
202 comp_cls = plugin.filter_component_classes['trimmer']
203 return self._graph.add_component(comp_cls, name, params)
204
88fdcc33 205 def _get_unique_comp_name(self, comp_spec):
d34e69cf 206 name = '{}-{}'.format(comp_spec.plugin_name,
c88be1c8 207 comp_spec.class_name)
88fdcc33
PP
208 comps_and_specs = itertools.chain(self._src_comps_and_specs,
209 self._flt_comps_and_specs)
d34e69cf 210
88fdcc33 211 if name in [comp_and_spec.comp.name for comp_and_spec in comps_and_specs]:
d34e69cf
PP
212 name += '-{}'.format(self._next_suffix)
213 self._next_suffix += 1
214
215 return name
216
88fdcc33 217 def _create_comp(self, comp_spec, comp_cls_type):
d34e69cf
PP
218 plugin = bt2.find_plugin(comp_spec.plugin_name)
219
220 if plugin is None:
221 raise bt2.Error('no such plugin: {}'.format(comp_spec.plugin_name))
222
88fdcc33
PP
223 if comp_cls_type == _CompClsType.SOURCE:
224 comp_classes = plugin.source_component_classes
225 else:
226 comp_classes = plugin.filter_component_classes
227
c88be1c8 228 if comp_spec.class_name not in comp_classes:
88fdcc33
PP
229 cc_type = 'source' if comp_cls_type == _CompClsType.SOURCE else 'filter'
230 raise bt2.Error('no such {} component class in "{}" plugin: {}'.format(cc_type,
231 comp_spec.plugin_name,
c88be1c8 232 comp_spec.class_name))
d34e69cf 233
c88be1c8 234 comp_cls = comp_classes[comp_spec.class_name]
88fdcc33 235 name = self._get_unique_comp_name(comp_spec)
cc81b5ab
PP
236 comp = self._graph.add_component(comp_cls, name, comp_spec.params,
237 comp_spec.logging_level)
d34e69cf
PP
238 return comp
239
240 def _get_free_muxer_input_port(self):
241 for port in self._muxer_comp.input_ports.values():
242 if not port.is_connected:
243 return port
244
da35796c 245 def _connect_src_comp_port(self, component, port):
d34e69cf
PP
246 # if this trace collection iterator is in stream intersection
247 # mode, we need this connection:
248 #
249 # port -> trimmer -> muxer
250 #
251 # otherwise, simply:
252 #
253 # port -> muxer
254 if self._stream_intersection_mode:
da35796c 255 trimmer_comp = self._create_stream_intersection_trimmer(component, port)
d34e69cf
PP
256 self._graph.connect_ports(port, trimmer_comp.input_ports['in'])
257 port_to_muxer = trimmer_comp.output_ports['out']
258 else:
259 port_to_muxer = port
260
261 self._graph.connect_ports(port_to_muxer, self._get_free_muxer_input_port())
262
da35796c 263 def _graph_port_added(self, component, port):
d34e69cf
PP
264 if not self._connect_ports:
265 return
266
da35796c 267 if type(port) is bt2.port._InputPort:
d34e69cf
PP
268 return
269
da35796c 270 if component not in [comp.comp for comp in self._src_comps_and_specs]:
d34e69cf
PP
271 # do not care about non-source components (muxer, trimmer, etc.)
272 return
273
da35796c 274 self._connect_src_comp_port(component, port)
d34e69cf
PP
275
276 def _build_graph(self):
277 self._graph = bt2.Graph()
da35796c 278 self._graph.add_port_added_listener(self._graph_port_added)
d34e69cf
PP
279 self._muxer_comp = self._create_muxer()
280
281 if self._begin_ns is not None or self._end_ns is not None:
282 trimmer_comp = self._create_trimmer(self._begin_ns,
283 self._end_ns, 'trimmer')
284 self._graph.connect_ports(self._muxer_comp.output_ports['out'],
285 trimmer_comp.input_ports['in'])
fa4c33e3 286 msg_iter_port = trimmer_comp.output_ports['out']
d34e69cf 287 else:
fa4c33e3 288 msg_iter_port = self._muxer_comp.output_ports['out']
d34e69cf 289
88fdcc33
PP
290 # create extra filter components (chained)
291 for comp_spec in self._flt_comp_specs:
292 comp = self._create_comp(comp_spec, _CompClsType.FILTER)
293 self._flt_comps_and_specs.append(_ComponentAndSpec(comp, comp_spec))
294
295 # connect the extra filter chain
296 for comp_and_spec in self._flt_comps_and_specs:
297 in_port = list(comp_and_spec.comp.input_ports.values())[0]
298 out_port = list(comp_and_spec.comp.output_ports.values())[0]
fa4c33e3
SM
299 self._graph.connect_ports(msg_iter_port, in_port)
300 msg_iter_port = out_port
88fdcc33 301
d34e69cf
PP
302 # Here we create the components, self._graph_port_added() is
303 # called when they add ports, but the callback returns early
304 # because self._connect_ports is False. This is because the
305 # self._graph_port_added() could not find the associated source
306 # component specification in self._src_comps_and_specs because
307 # it does not exist yet (it needs the created component to
308 # exist).
309 for comp_spec in self._src_comp_specs:
88fdcc33
PP
310 comp = self._create_comp(comp_spec, _CompClsType.SOURCE)
311 self._src_comps_and_specs.append(_ComponentAndSpec(comp, comp_spec))
d34e69cf
PP
312
313 # Now we connect the ports which exist at this point. We allow
314 # self._graph_port_added() to automatically connect _new_ ports.
315 self._connect_ports = True
316
317 for comp_and_spec in self._src_comps_and_specs:
318 # Keep a separate list because comp_and_spec.output_ports
319 # could change during the connection of one of its ports.
320 # Any new port is handled by self._graph_port_added().
321 out_ports = [port for port in comp_and_spec.comp.output_ports.values()]
322
323 for out_port in out_ports:
da35796c 324 if out_port.is_connected:
d34e69cf
PP
325 continue
326
da35796c 327 self._connect_src_comp_port(comp_and_spec.comp, out_port)
d34e69cf 328
fa4c33e3 329 # create this trace collection iterator's message iterator
da35796c 330 self._msg_iter = self._graph.create_output_port_message_iterator(msg_iter_port)
This page took 0.047429 seconds and 4 git commands to generate.