From 3aa456fd4b2fdf734b03e05232deb63568e61641 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Wed, 7 Jun 2017 20:48:36 -0500 Subject: [PATCH] Move reset_auto_ids to mixin module and add set_auto_id and reserve_ids functions --- openmc/geometry.py | 8 -------- openmc/mixin.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/openmc/geometry.py b/openmc/geometry.py index 6b1e6cb72..0b02210b7 100644 --- a/openmc/geometry.py +++ b/openmc/geometry.py @@ -7,14 +7,6 @@ from six import string_types import openmc from openmc.clean_xml import sort_xml_elements, clean_xml_indentation from openmc.checkvalue import check_type -from openmc.mixin import IDManagerMixin - - -def reset_auto_ids(): - """Reset counters for all auto-generated IDs""" - for cls in IDManagerMixin.__subclasses__(): - cls.used_ids.clear() - cls.next_id = 1 class Geometry(object): diff --git a/openmc/mixin.py b/openmc/mixin.py index 0b0c5f519..555c3b158 100644 --- a/openmc/mixin.py +++ b/openmc/mixin.py @@ -57,3 +57,39 @@ class IDManagerMixin(object): else: cls.used_ids.add(uid) self._id = uid + + +def reset_auto_ids(): + """Reset counters for all auto-generated IDs""" + for cls in IDManagerMixin.__subclasses__(): + cls.used_ids.clear() + cls.next_id = 1 + + +def reserve_ids(ids, cls=None): + """Reserve a set of IDs that won't be used for auto-generated IDs. + + Parameters + ---------- + ids : iterable of int + IDs to reserve + + """ + if cls is None: + for cls in IDManagerMixin.__subclasses__(): + cls.used_ids |= set(ids) + else: + cls.used_ids |= set(ids) + + +def set_auto_id(next_id): + """Set the next ID for auto-generated IDs. + + Parameters + ---------- + next_id : int + The next ID to assign to objects with auto-generated IDs. + + """ + for cls in IDManagerMixin.__subclasses__(): + cls.next_id = next_id