From c0f7459ff39c5735dd873cfb658a8885b363a603 Mon Sep 17 00:00:00 2001
From: Mikolaj Adam Kowalski <mak60@cam.ac.uk>
Date: Mon, 2 Mar 2020 19:41:56 +0000
Subject: [PATCH] Implement getter for Cell.atoms

Make type checking in Cell.volume more compact
---
 openmc/cell.py | 36 +++++++++++++++++++++++++++++++-----
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/openmc/cell.py b/openmc/cell.py
index 792f152bef..3ecab7ba19 100644
--- a/openmc/cell.py
+++ b/openmc/cell.py
@@ -184,6 +184,35 @@ class Cell(IDManagerMixin):
     def volume(self):
         return self._volume
 
+    @property
+    def atoms(self):
+        if self._atoms is None:
+            if self._volume is None:
+                msg = 'Cannot calculate atoms content becouse no volume '\
+                      'is set. Use Cell.volume to provide it or perform '\
+                      'stochastic volume calculation.'
+                raise ValueError(msg)
+
+            elif self._fill is None:
+                msg = 'Cell is filled with void. It contains no atoms.'
+                raise ValueError(msg)
+
+            elif isinstance(self._fill, (openmc.Universe, openmc.Lattice)):
+                msg = 'Universe and Lattice cells can contain multiple '\
+                      'materials. Atoms content must be calculated with '\
+                      'stochastic volume calculation'
+                raise ValueError(msg)
+
+            elif isinstance(self._fill, openmc.Material):
+                # Get atomic Densities
+                self._atoms = self._fill.get_nuclide_atom_densities()
+
+                # Convert to total number of atoms
+                for key, nuclide in self._atoms.items():
+                    self._atoms[key] = (nuclide[0], nuclide[1] * self._volume)
+
+        return self._atoms
+
     @property
     def paths(self):
         if self._paths is None:
@@ -287,14 +316,11 @@ class Cell(IDManagerMixin):
     @volume.setter
     def volume(self, volume):
         if volume is not None:
-            try:
-                cv.check_type('cell volume', volume, Real)
-            except TypeError:
-                cv.check_type('cell volume', volume, UFloat)
+            cv.check_type('cell volume', volume, (Real, UFloat))
             cv.check_greater_than('cell volume', volume, 0.0)
         self._volume = volume
 
-        # Forget now invalid info about atoms content
+        # Info about atoms content can now be invalid
         # (sice volume has just changed)
         if self._atoms is not None:
             self._atoms = None