Repository: CORE-GATECH-GROUP/serpent-tools
Pull request #309 (2019-05-08 20:45:39)

API Add setitem/getitem methods for HomogUniv object
Adds the ability to treat `HomogUniv` objects like a dictionary, at least in regards to accessing expected values of group constants. With this change, one can do the following
```
>>> h = HomogUniv("setget", None, None, None)
>>> h['infDummy'] = 1.0
>>> h.infExp
{'infDummy': 1.0}
>>> h.infUnc:
{}
>>> h['infDummy']
1.0
```
For obtaining uncertainties, the best way is still the [`get` method](https://serpent-tools.readthedocs.io/en/0.7.0/containers/universes.html#serpentTools.objects.HomogUniv.get). Tests have been added to verify we are setting data correctly, and in the right place, and can retrieve it.

Did some documentation cleanup, and cleaned up the interior lookup method used to obtain the right dictionary to place data and to retrieve data.


Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index c9a0dd3..307d9f7 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -6,6 +6,13 @@
 Changelog
 =========
 
+Next
+====
+
+* Add :meth:`~serpentTools.objects.HomogUniv.__getitem__` and 
+  :meth:`~serpentTools.objects.HomogUniv.__setitem__` convenience
+  methods for accessing expected values on |HomogUniv| objects
+
 .. _v0.7.0:
 
 :release-tag:`0.7.0`
diff --git a/docs/containers/universes.rst b/docs/containers/universes.rst
index 3835d37..a9e61f3 100644
--- a/docs/containers/universes.rst
+++ b/docs/containers/universes.rst
@@ -4,3 +4,4 @@ Homogenized Universes
 =====================
 
 .. autoclass:: serpentTools.objects.HomogUniv
+    :special-members: __getitem__, __setitem__, __bool__, __nonzero__
diff --git a/serpentTools/objects/containers.py b/serpentTools/objects/containers.py
index fc2d24e..23bd7e4 100644
--- a/serpentTools/objects/containers.py
+++ b/serpentTools/objects/containers.py
@@ -59,13 +59,6 @@ __all__ = (
 CRIT_RE = compile('K[EI][NF]F')
 
 
-def isNonNeg(value):
-    """Return true if a value is None or non-negative"""
-    if value is None:
-        return True
-    return value >= 0
-
-
 class HomogUniv(NamedObject):
     """
     Class for storing homogenized universe specifications and retrieving them
@@ -108,23 +101,23 @@ class HomogUniv(NamedObject):
     reshaped: bool
         ``True`` if scattering matrices have been reshaped to square
         matrices. Otherwise, these matrices are stored as vectors.
-    groups: None or :py:class:`numpy.array`
+    groups: None or :class:`numpy.array`
         Group boundaries from highest to lowest
     numGroups: None or int
         Number of energy groups bounded by ``groups``
-    microGroups: None or :py:class:`numpy.array`
+    microGroups: None or :class:`numpy.array`
         Micro group structure used to produce group constants.
         Listed from lowest to highest
 
     Raises
     ------
-    :class:`~serpentTools.SerpentToolsException`
+    :class:`serpentTools.SerpentToolsException`
         If a negative value of burnup, step, or burnup days is passed
 
     """
 
     def __init__(self, name, bu, step, day):
-        if not all(isNonNeg(xx) for xx in (bu, step, day)):
+        if not all(xx is None or xx >= 0 for xx in (bu, step, day)):
             tail = ['{}: {}'.format(valName, val)
                     for valName, val in zip(('burnup', 'index', 'days'),
                                             (bu, step, day))]
@@ -286,6 +279,10 @@ class HomogUniv(NamedObject):
             If the variable requested is not stored on the
             object
 
+        See Also
+        --------
+        :meth:`__getitem__` to directly access data witout uncertainties
+
         """
         # 1. Check the input values
         if not isinstance(uncertainty, bool):
@@ -306,18 +303,41 @@ class HomogUniv(NamedObject):
             raise KeyError(
                 "Variable {} absent from uncertainty dictionary".format(
                     variableName))
-        dx = setter.get(variableName)
-        return x, dx
+
+        return x, setter[variableName]
+
+    def __getitem__(self, gcname):
+        """
+        Return just the group constant with this name, no uncertainty
+
+        To return data with uncertainties, or to return uncertainties,
+        use :meth:`get`.
+        """
+        return self._lookup(gcname, False)[gcname]
+
+    def __setitem__(self, gckey, gcvalue):
+        """
+        Set the expected value of gckey to be gcvalue
+
+        No conversions are placed on the variable name. What you
+        pass is what gets set.
+
+        For uncertainties, or to convert variable
+        names to ``mixedCase``, use :meth:`addData`.
+
+        Much like a dictionary, this will overwrite existing data.
+        """
+        self._lookup(gckey, False)[gckey] = gcvalue
 
     def _lookup(self, variableName, uncertainty):
         if 'inf' == variableName[:3]:
-            if not uncertainty:
-                return self.infExp
-            return self.infUnc
+            if uncertainty:
+                return self.infUnc
+            return self.infExp
         elif "b1" == variableName[:2]:
-            if not uncertainty:
-                return self.b1Exp
-            return self.b1Unc
+            if uncertainty:
+                return self.b1Unc
+            return self.b1Exp
         return self.gcUnc if uncertainty else self.gc
 
     @magicPlotDocDecorator
@@ -383,13 +403,9 @@ class HomogUniv(NamedObject):
                 raise IndexError(
                     "Need equal number of labels for plot quantities. "
                     "Given {} expected: {}".format(len(labels), len(qtys)))
+
         for key, label in zip(qtys, labels):
-            valD = self._lookup(key, False)
-            if key not in valD:
-                warning("{} not found on object. Will not plot."
-                        .format(key))
-                continue
-            yVals = valD[key]
+            yVals = self.__getitem__(key)
             if len(yVals.shape) != 1 and 1 not in yVals.shape:
                 warning("Data for {} is not 1D. Will not plot"
                         .format(key))
@@ -685,7 +701,7 @@ class BranchContainer(BaseObject):
 
         Raises
         ------
-        :class:`~serpentTools.SerpentToolsException`
+        :class:`serpentTools.SerpentToolsException`
             If passed a value of ``burnDays`` and set up to work with burnup,
             or vice versa
         """
@@ -733,9 +749,9 @@ class BranchContainer(BaseObject):
 
         Raises
         ------
-        KeyError:
+        KeyError
             If the requested universe could not be found
-        :class:`~serpentTools.SerpentToolsException`
+        :class:`serpentTools.SerpentToolsException`
             If neither burnup nor index are given
         """
         if burnup is None and index is None:
diff --git a/serpentTools/objects/xsdata.py b/serpentTools/objects/xsdata.py
index cf690c3..a1c1890 100644
--- a/serpentTools/objects/xsdata.py
+++ b/serpentTools/objects/xsdata.py
@@ -145,9 +145,9 @@ class XSData(NamedObject):
 
         Raises
         ------
-        ModuleNotFoundError:
-            if you don't have pandas.
-        TypeError:
+        ImportError
+            If ``pandas`` is not installed
+        TypeError
             if MT numbers that don't make sense come up
         """
         import pandas as pd
@@ -226,7 +226,7 @@ class XSData(NamedObject):
 
         Raises
         ------
-        TypeError:
+        TypeError
             if MT numbers that don't make sense come up
         """
 


Test patch:
diff --git a/serpentTools/tests/test_container.py b/serpentTools/tests/test_container.py
index 61accb6..35de971 100644
--- a/serpentTools/tests/test_container.py
+++ b/serpentTools/tests/test_container.py
@@ -1,6 +1,6 @@
 """Test the container object. """
 
-import unittest
+from unittest import TestCase
 from itertools import product
 
 from six import iteritems
@@ -8,12 +8,13 @@ from numpy import arange, ndarray, float64
 from numpy.testing import assert_array_equal
 
 from serpentTools.objects.containers import HomogUniv
+from serpentTools.messages import SerpentToolsException
 from serpentTools.tests import compareDictOfArrays
 
 NUM_GROUPS = 5
 
 
-class _HomogUnivTestHelper(unittest.TestCase):
+class _HomogUnivTestHelper(TestCase):
     """Class that runs the tests for the two sub-classes
 
     Subclasses will differ in how the ``mat`` data
@@ -129,6 +130,36 @@ class ReshapedHomogUnivTester(_HomogUnivTestHelper):
         return univ, vec, mat.reshape(NUM_GROUPS, NUM_GROUPS, order="F")
 
 
+class SimpleHomogUnivTester(TestCase):
+    """Class for simple tests on HomogUniv objects"""
+
+    def _wrapBadInit(self, *initArgs):
+        msg = "Args: {}".format(initArgs)
+        with self.assertRaises(SerpentToolsException, msg=msg):
+            HomogUniv('badInit', *initArgs)
+
+    def test_badinit(self):
+        """Verify that a HomogUniv cannot be constructed with negative time"""
+        self._wrapBadInit(-1, 0, None)
+        self._wrapBadInit(None, None, -10.0)
+        # good init
+        HomogUniv('good', None, None, None)
+
+    def test_setGet(self):
+        infKey = 'infTestData'
+        infVal = arange(2)
+        univ = HomogUniv('setGetTest', None, None, None)
+        self.assertEqual(0, len(univ.infExp), msg='Before set')
+        self.assertEqual(0, len(univ.infUnc), msg='Before set')
+        univ[infKey] = infVal
+        self.assertEqual(1, len(univ.infExp), msg='After set')
+        self.assertTrue(infKey in univ.infExp)
+        self.assertEqual(0, len(univ.infUnc), msg='After set')
+        self.assertIs(univ[infKey], infVal)
+        with self.assertRaises(KeyError):
+            univ['this should not exist']
+
+
 def getParams():
     """Return the universe, vector, and matrix for testing."""
     univ = HomogUniv(300, 0, 0, 0)
@@ -140,7 +171,7 @@ def getParams():
 del _HomogUnivTestHelper
 
 
-class UnivTruthTester(unittest.TestCase):
+class UnivTruthTester(TestCase):
     """Class that tests the various boolean evaluations for HomogUniv"""
 
     def setUp(self):
@@ -169,7 +200,7 @@ class UnivTruthTester(unittest.TestCase):
         self.assertTrue(univ.hasData, msg=msg)
 
 
-class HomogUnivIntGroupsTester(unittest.TestCase):
+class HomogUnivIntGroupsTester(TestCase):
     """Class that ensures number of groups is stored as ints."""
 
     def setUp(self):
@@ -203,4 +234,5 @@ class HomogUnivIntGroupsTester(unittest.TestCase):
 
 
 if __name__ == '__main__':
-    unittest.main()
+    from unittest import main
+    main()