diff --git a/cbutil/util.py b/cbutil/util.py index 292c5fca55a80b0f154b5cf25e67317e456b002f..60daeacc07333333029d7fd1b43c0d68568dd91a 100644 --- a/cbutil/util.py +++ b/cbutil/util.py @@ -2,17 +2,38 @@ import logging import operator import os import time +import re from datetime import datetime from functools import reduce from pathlib import Path -from typing import Tuple, Union +from typing import Tuple, Union, List import git logger = logging.getLogger(__file__) +def parse_option_int_array(specification: str): + """ Parses the specification from Petsc PetscOptionGetIntarray. + """ + elements = specification.split(',') + for element in elements: + if (match := re.match(r"(\d+)-(\d+)(:(\d+))?", element)): + lower = int(match.group(1)) + upper = int(match.group(2)) + step = 1 + if match.lastindex == 3: + step = int(match.group(4)) + yield from range(lower, upper, step) + else: + yield int(element) + + +def get_option_int_array(specification: str) -> List: + return sorted(set(parse_option_int_array(specification))) + + def remove_none(old_dict: dict) -> dict: """Remove all entries with None as value. """ return {k: v for k, v in old_dict.items() if v is not None} diff --git a/tests/test_util.py b/tests/test_util.py index d201a752877830dddd780a7c1cd232060904ec7a..c8642fd1b97f79d1e8d8e26bf73e2541bcd59276 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -1,4 +1,4 @@ -from cbutil.util import read_file_line_wise, get_git_infos, time_conversion +from cbutil.util import read_file_line_wise, get_git_infos, time_conversion, parse_option_int_array from contextlib import contextmanager from pathlib import Path import os @@ -39,3 +39,11 @@ def test_time_conv(): time_conversion("ajsdkfljaskldj") with pytest.raises(ValueError): time_conversion("2022-08-0315:34:46") + + +def test_parse_option_int_array(): + assert list(parse_option_int_array("0,1")) == [0, 1] + assert list(parse_option_int_array("0-10")) == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + assert list(parse_option_int_array("0-1")) == [0] + assert list(parse_option_int_array("0-10:2")) == [0, 2, 4, 6, 8] + assert list(parse_option_int_array("0-10:2,11-13,20")) == [0, 2, 4, 6, 8, 11, 12, 20]