Fix capi.Cell.__new__

openmc/capi/cell.py

@@ -7,12 +7,15 @@ from numpy.ctypeslib import as_array
 
 from . import _dll
 from .core import _FortranObjectWithID
-from .error import _error_handler
+from .error import _error_handler, AllocationError, InvalidIDError
 from .material import Material
 
 __all__ = ['Cell', 'cells']
 
 # Cell functions
+_dll.openmc_extend_cells.argtypes = [c_int32, POINTER(c_int32), POINTER(c_int32)]
+_dll.openmc_extend_cells.restype = c_int
+_dll.openmc_extend_cells.errcheck = _error_handler
 _dll.openmc_cell_get_id.argtypes = [c_int32, POINTER(c_int32)]
 _dll.openmc_cell_get_id.restype = c_int
 _dll.openmc_cell_get_id.errcheck = _error_handler
@@ -51,11 +54,35 @@ class Cell(_FortranObjectWithID):
     """
     __instances = WeakValueDictionary()
 
-    def __new__(cls, *args):
-        if args not in cls.__instances:
-            instance = super(Cell, self).__new__(cls)
-            cls.__instances[args] = instance
-        return cls.__instances[args]
+    def __new__(cls, uid=None, new=True, index=None):
+        mapping = cells
+        if index is None:
+            if new:
+                # Determine ID to assign
+                if uid is None:
+                    try:
+                        uid = max(mapping) + 1
+                    except ValueError:
+                        uid = 1
+                else:
+                    if uid in mapping:
+                        raise AllocationError('A cell with ID={} has already '
+                                              'been allocated.'.format(uid))
+
+                index = c_int32()
+                _dll.openmc_extend_cells(1, index, None)
+                index = index.value
+            else:
+                index = mapping[uid]._index
+
+        if index not in cls.__instances:
+            instance = super(Cell, cls).__new__(cls)
+            instance._index = index
+            if uid is not None:
+                instance.id = uid
+            cls.__instances[index] = instance
+
+        return cls.__instances[index]
 
     @property
     def id(self):
@@ -113,11 +140,11 @@ class _CellMapping(Mapping):
         except (AllocationError, InvalidIDError) as e:
             # __contains__ expects a KeyError to work correctly
             raise KeyError(str(e))
-        return Cell(index.value)
+        return Cell(index=index.value)
 
     def __iter__(self):
         for i in range(len(self)):
-            yield Cell(i + 1).id
+            yield Cell(index=i + 1).id
 
     def __len__(self):
         return c_int32.in_dll(_dll, 'n_cells').value

src/api.F90

@@ -40,6 +40,7 @@ module openmc_api
   public :: openmc_energy_filter_get_bins
   public :: openmc_energy_filter_set_bins
   public :: openmc_extend_filters
+  public :: openmc_extend_cells
   public :: openmc_extend_materials
   public :: openmc_extend_tallies
   public :: openmc_filter_get_id

src/geometry_header.F90

@@ -427,6 +427,38 @@ contains
 !                               C API FUNCTIONS
 !===============================================================================
 
+  function openmc_extend_cells(n, index_start, index_end) result(err) bind(C)
+    ! Extend the cells array by n elements
+    integer(C_INT32_T), value, intent(in) :: n
+    integer(C_INT32_T), optional, intent(out) :: index_start
+    integer(C_INT32_T), optional, intent(out) :: index_end
+    integer(C_INT) :: err
+
+    type(Cell), allocatable :: temp(:) ! temporary cells array
+
+    if (n_cells == 0) then
+      ! Allocate cells array
+      allocate(cells(n))
+    else
+      ! Allocate cells array with increased size
+      allocate(temp(n_cells + n))
+
+      ! Copy original cells to temporary array
+      temp(1:n_cells) = cells
+
+      ! Move allocation from temporary array
+      call move_alloc(FROM=temp, TO=cells)
+    end if
+
+    ! Return indices in cells array
+    if (present(index_start)) index_start = n_cells + 1
+    if (present(index_end)) index_end = n_cells + n
+    n_cells = n_cells + n
+
+    err = 0
+  end function openmc_extend_cells
+
+
   function openmc_get_cell_index(id, index) result(err) bind(C)
     ! Return the index in the cells array of a cell with a given ID
     integer(C_INT32_T), value :: id