diff --git a/intervals/__init__.py b/intervals/__init__.py index 2cc88b6..4c949ad 100644 --- a/intervals/__init__.py +++ b/intervals/__init__.py @@ -72,6 +72,21 @@ def wrapper(self, arg): return wrapper +class ClosedInterval(type): + """ + Supports initialization of intervals using square brackets and makes them + closed intervals. + + eg. + + IntInterval[1, 4] == IntInterval([1, 4]) + """ + def __getitem__(self, bounds): + lower_inc = upper_inc = True + return self(bounds, lower_inc, upper_inc) + + +@six.add_metaclass(ClosedInterval) @total_ordering class AbstractInterval(object): step = None @@ -151,6 +166,20 @@ def __init__( 30 """ + + # This if-block adds support for parentheses as open intervals. + # Note: If the interval is initialized with the parentheses with two + # objects of same type, eg. + # IntInterval(1, 4) + # the bounds and lower_inc are received of that type and + # upper_inc is None. + # + # eg. + # IntInterval(1, 4) == IntInterval((1, 4)) + if type(bounds) == type(lower_inc) and not upper_inc: + bounds = (bounds, lower_inc) + lower_inc = upper_inc = None + self.lower, self.upper, self.lower_inc, self.upper_inc = ( self.parser(bounds, lower_inc, upper_inc) ) diff --git a/tests/test_initialization.py b/tests/test_initialization.py index 6840bb8..eaac817 100644 --- a/tests/test_initialization.py +++ b/tests/test_initialization.py @@ -88,6 +88,12 @@ def test_supports_integers(self): assert interval.lower_inc assert interval.upper_inc + def test_uses_two_numbers_with_parentheses_as_open_interval(self): + assert IntInterval(1, 2) == IntInterval((1, 2)) + + def test_uses_two_numbers_with_square_brackets_as_closed_interval(self): + assert IntInterval[1, 2] == IntInterval([1, 2]) + @mark.parametrize('number_range', ( (3, 2),