diff --git a/pytools/__init__.py b/pytools/__init__.py index 7bc99bb919c8020ac6eeec0ae059c3184cfbf5eb..a9c58ce26e893daed48c17d6ca894080ceb82ac8 100644 --- a/pytools/__init__.py +++ b/pytools/__init__.py @@ -1855,10 +1855,9 @@ def generate_unique_names(prefix): def generate_numbered_unique_names(prefix, num=None): - orig_num = num - num = 0 - if orig_num is None: - yield (num, prefix) + if num is None: + yield (0, prefix) + num = 0 while True: name = "%s_%d" % (prefix, num) diff --git a/test/test_pytools.py b/test/test_pytools.py index d73553d97c639516a937f56786ffd3840a2519b2..5346c25591d142a5af4c3d677ec8b92d5f03bbe5 100644 --- a/test/test_pytools.py +++ b/test/test_pytools.py @@ -195,6 +195,17 @@ def test_diskdict(): del d +def test_generate_numbered_unique_names(): + from pytools import generate_numbered_unique_names + + gen = generate_numbered_unique_names("a") + assert next(gen) == (0, "a") + assert next(gen) == (1, "a_0") + + gen = generate_numbered_unique_names("b", 6) + assert next(gen) == (7, "b_6") + + if __name__ == "__main__": if len(sys.argv) > 1: exec(sys.argv[1])