bt2: _IntegerRangeSet.add(): make it possible to add a single value

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

diff --git a/src/bindings/python/bt2/bt2/integer_range_set.py b/src/bindings/python/bt2/bt2/integer_range_set.py
index 2e79da04a3d3f39d2dd772a04bed6b73600191a5..09d7b08fcdce95fca6fe1ea059247da4f84f2f1d 100644 (file)
--- a/src/bindings/python/bt2/bt2/integer_range_set.py
+++ b/src/bindings/python/bt2/bt2/integer_range_set.py
@@ -26,12 +26,20 @@ import bt2
 
 
 class _IntegerRange:
-    def __init__(self, lower, upper):
+    def __init__(self, lower, upper=None):
         self._check_type(lower)
+
+        if upper is None:
+            upper = lower
+
         self._check_type(upper)
 
         if lower > upper:
-            raise ValueError("range's lower bound ({}) is greater than its upper bound ({})".format(lower, upper))
+            raise ValueError(
+                "range's lower bound ({}) is greater than its upper bound ({})".format(
+                    lower, upper
+                )
+            )
 
         self._lower = lower
         self._upper = upper
@@ -56,10 +64,12 @@ class _IntegerRange:
 
 
 class SignedIntegerRange(_IntegerRange):
+    _is_type = staticmethod(utils._is_int64)
     _check_type = staticmethod(utils._check_int64)
 
 
 class UnsignedIntegerRange(_IntegerRange):
+    _is_type = staticmethod(utils._is_uint64)
     _check_type = staticmethod(utils._check_uint64)
 
 
@@ -68,7 +78,7 @@ class _IntegerRangeSet(object._SharedObject, collections.abc.MutableSet):
         ptr = self._create_range_set()
 
         if ptr is None:
-            raise bt2.CreationError('cannot create range set object')
+            raise bt2._MemoryError('cannot create range set object')
 
         super().__init__(ptr)
 
@@ -113,12 +123,14 @@ class _IntegerRangeSet(object._SharedObject, collections.abc.MutableSet):
 
     def add(self, rg):
         if type(rg) is not self._range_type:
-            # assume it's a simple pair (will raise if it's not)
-            rg = self._range_type(rg[0], rg[1])
+            if self._range_type._is_type(rg):
+                rg = self._range_type(rg)
+            else:
+                # assume it's a simple pair (will raise if it's not)
+                rg = self._range_type(rg[0], rg[1])
 
         status = self._add_range(self._ptr, rg.lower, rg.upper)
-        utils._handle_func_status(status,
-                                  'cannot add range to range set object')
+        utils._handle_func_status(status, 'cannot add range to range set object')
 
     def discard(self, rg):
         raise NotImplementedError
@@ -127,9 +139,13 @@ class _IntegerRangeSet(object._SharedObject, collections.abc.MutableSet):
 class SignedIntegerRangeSet(_IntegerRangeSet):
     _get_ref = staticmethod(native_bt.integer_range_set_signed_get_ref)
     _put_ref = staticmethod(native_bt.integer_range_set_signed_put_ref)
-    _as_range_set_ptr = staticmethod(native_bt.integer_range_set_signed_as_range_set_const)
+    _as_range_set_ptr = staticmethod(
+        native_bt.integer_range_set_signed_as_range_set_const
+    )
     _create_range_set = staticmethod(native_bt.integer_range_set_signed_create)
-    _borrow_range_by_index_ptr = staticmethod(native_bt.integer_range_set_signed_borrow_range_by_index_const)
+    _borrow_range_by_index_ptr = staticmethod(
+        native_bt.integer_range_set_signed_borrow_range_by_index_const
+    )
     _range_get_lower = staticmethod(native_bt.integer_range_signed_get_lower)
     _range_get_upper = staticmethod(native_bt.integer_range_signed_get_upper)
     _add_range = staticmethod(native_bt.integer_range_set_signed_add_range)
@@ -140,9 +156,13 @@ class SignedIntegerRangeSet(_IntegerRangeSet):
 class UnsignedIntegerRangeSet(_IntegerRangeSet):
     _get_ref = staticmethod(native_bt.integer_range_set_unsigned_get_ref)
     _put_ref = staticmethod(native_bt.integer_range_set_unsigned_put_ref)
-    _as_range_set_ptr = staticmethod(native_bt.integer_range_set_unsigned_as_range_set_const)
+    _as_range_set_ptr = staticmethod(
+        native_bt.integer_range_set_unsigned_as_range_set_const
+    )
     _create_range_set = staticmethod(native_bt.integer_range_set_unsigned_create)
-    _borrow_range_by_index_ptr = staticmethod(native_bt.integer_range_set_unsigned_borrow_range_by_index_const)
+    _borrow_range_by_index_ptr = staticmethod(
+        native_bt.integer_range_set_unsigned_borrow_range_by_index_const
+    )
     _range_get_lower = staticmethod(native_bt.integer_range_unsigned_get_lower)
     _range_get_upper = staticmethod(native_bt.integer_range_unsigned_get_upper)
     _add_range = staticmethod(native_bt.integer_range_set_unsigned_add_range)