bt2: pass custom Python object to Python component's __init__()

[babeltrace.git] / src / bindings / python / bt2 / bt2 / trace_collection_message_iterator.py

diff --git a/src/bindings/python/bt2/bt2/trace_collection_message_iterator.py b/src/bindings/python/bt2/bt2/trace_collection_message_iterator.py
index c6d147afb743a04477e555b74fe110cbaa0ec07d..4e7361e7ee877e4df02b1e5e005f20c74df3cc25 100644 (file)
--- a/src/bindings/python/bt2/bt2/trace_collection_message_iterator.py
+++ b/src/bindings/python/bt2/bt2/trace_collection_message_iterator.py
@@ -23,7 +23,9 @@
 from bt2 import utils
 import bt2
 import itertools
-import bt2.message_iterator
+from bt2 import message_iterator as bt2_message_iterator
+from bt2 import logging as bt2_logging
+from bt2 import port as bt2_port
 import datetime
 from collections import namedtuple
 import numbers
@@ -39,7 +41,8 @@ class ComponentSpec:
         plugin_name,
         class_name,
         params=None,
-        logging_level=bt2.logging.LoggingLevel.NONE,
+        obj=None,
+        logging_level=bt2_logging.LoggingLevel.NONE,
     ):
         utils._check_str(plugin_name)
         utils._check_str(class_name)
@@ -47,6 +50,7 @@ class ComponentSpec:
         self._plugin_name = plugin_name
         self._class_name = class_name
         self._logging_level = logging_level
+        self._obj = obj
 
         if type(params) is str:
             self._params = bt2.create_value({'inputs': [params]})
@@ -69,6 +73,10 @@ class ComponentSpec:
     def params(self):
         return self._params
 
+    @property
+    def obj(self):
+        return self._obj
+
 
 # datetime.datetime or integral to nanoseconds
 def _get_ns(obj):
@@ -94,7 +102,7 @@ class _CompClsType:
     FILTER = 1
 
 
-class TraceCollectionMessageIterator(bt2.message_iterator._MessageIterator):
+class TraceCollectionMessageIterator(bt2_message_iterator._MessageIterator):
     def __init__(
         self,
         source_component_specs,
@@ -149,19 +157,19 @@ class TraceCollectionMessageIterator(bt2.message_iterator._MessageIterator):
         try:
             inputs = src_comp_and_spec.spec.params['inputs']
         except Exception as e:
-            raise bt2.Error(
+            raise ValueError(
                 'all source components must be created with an "inputs" parameter in stream intersection mode'
             ) from e
 
         params = {'inputs': inputs}
 
-        # query the port's component for the `trace-info` object which
-        # contains the stream intersection range for each exposed
-        # trace
-        query_exec = bt2.QueryExecutor()
-        trace_info_res = query_exec.query(
-            src_comp_and_spec.comp.cls, 'trace-info', params
+        # query the port's component for the `babeltrace.trace-info`
+        # object which contains the stream intersection range for each
+        # exposed trace
+        query_exec = bt2.QueryExecutor(
+            src_comp_and_spec.comp.cls, 'babeltrace.trace-info', params
         )
+        trace_info_res = query_exec.query()
         begin = None
         end = None
 
@@ -178,7 +186,7 @@ class TraceCollectionMessageIterator(bt2.message_iterator._MessageIterator):
             pass
 
         if begin is None or end is None:
-            raise bt2.Error(
+            raise RuntimeError(
                 'cannot find stream intersection range for port "{}"'.format(port.name)
             )
 
@@ -189,10 +197,10 @@ class TraceCollectionMessageIterator(bt2.message_iterator._MessageIterator):
         plugin = bt2.find_plugin('utils')
 
         if plugin is None:
-            raise bt2.Error('cannot find "utils" plugin (needed for the muxer)')
+            raise RuntimeError('cannot find "utils" plugin (needed for the muxer)')
 
         if 'muxer' not in plugin.filter_component_classes:
-            raise bt2.Error(
+            raise RuntimeError(
                 'cannot find "muxer" filter component class in "utils" plugin'
             )
 
@@ -203,10 +211,10 @@ class TraceCollectionMessageIterator(bt2.message_iterator._MessageIterator):
         plugin = bt2.find_plugin('utils')
 
         if plugin is None:
-            raise bt2.Error('cannot find "utils" plugin (needed for the trimmer)')
+            raise RuntimeError('cannot find "utils" plugin (needed for the trimmer)')
 
         if 'trimmer' not in plugin.filter_component_classes:
-            raise bt2.Error(
+            raise RuntimeError(
                 'cannot find "trimmer" filter component class in "utils" plugin'
             )
 
@@ -242,7 +250,7 @@ class TraceCollectionMessageIterator(bt2.message_iterator._MessageIterator):
         plugin = bt2.find_plugin(comp_spec.plugin_name)
 
         if plugin is None:
-            raise bt2.Error('no such plugin: {}'.format(comp_spec.plugin_name))
+            raise ValueError('no such plugin: {}'.format(comp_spec.plugin_name))
 
         if comp_cls_type == _CompClsType.SOURCE:
             comp_classes = plugin.source_component_classes
@@ -251,7 +259,7 @@ class TraceCollectionMessageIterator(bt2.message_iterator._MessageIterator):
 
         if comp_spec.class_name not in comp_classes:
             cc_type = 'source' if comp_cls_type == _CompClsType.SOURCE else 'filter'
-            raise bt2.Error(
+            raise ValueError(
                 'no such {} component class in "{}" plugin: {}'.format(
                     cc_type, comp_spec.plugin_name, comp_spec.class_name
                 )
@@ -260,7 +268,7 @@ class TraceCollectionMessageIterator(bt2.message_iterator._MessageIterator):
         comp_cls = comp_classes[comp_spec.class_name]
         name = self._get_unique_comp_name(comp_spec)
         comp = self._graph.add_component(
-            comp_cls, name, comp_spec.params, comp_spec.logging_level
+            comp_cls, name, comp_spec.params, comp_spec.obj, comp_spec.logging_level
         )
         return comp
 
@@ -291,7 +299,7 @@ class TraceCollectionMessageIterator(bt2.message_iterator._MessageIterator):
         if not self._connect_ports:
             return
 
-        if type(port) is bt2.port._InputPort:
+        if type(port) is bt2_port._InputPort:
             return
 
         if component not in [comp.comp for comp in self._src_comps_and_specs]: