Provide better composite iteration methods

armi/bookkeeping/db/database.py

@@ -829,11 +829,11 @@ def loadReadOnly(self, cycle, node, statePointName=None):
         return r
 
     @staticmethod
-    def _setParamsBeforeFreezing(r):
+    def _setParamsBeforeFreezing(r: Reactor):
         """Set some special case parameters before they are made read-only."""
-        for child in r.getChildren(deep=True):
-            if not isinstance(child, Component):
-                continue
+        for child in r.iterChildren(
+            deep=True, predicate=lambda c: isinstance(c, Component)
+        ):
             # calling Component.getVolume() sets the volume parameter
             child.getVolume()
 

armi/bookkeeping/report/reportInterface.py

@@ -184,7 +184,7 @@ def reportSFP(sfp):
         runLog.important(title)
         runLog.important("-" * len(title))
         totFis = 0.0
-        for a in sfp.getChildren():
+        for a in sfp:
             runLog.important(
                 "{assembly:15s} discharged at t={dTime:10f} after {residence:10f} yrs. It entered at cycle: {cycle}. "
                 "It has {fiss:10f} kg (x {mult}) fissile and peak BU={bu:.2f} %.".format(
@@ -208,16 +208,16 @@ def reportSFP(sfp):
     @staticmethod
     def countAssembliesSFP(sfp):
         """Report on the count of assemblies in the SFP at each timestep."""
-        if not sfp.getChildren():
+        if not len(sfp):
             return
 
         runLog.important("Count:")
         totCount = 0
         thisTimeCount = 0
-        a = sfp.getChildren()[0]
+        a = sfp[0]
         lastTime = a.getAge() / units.DAYS_PER_YEAR + a.p.chargeTime
 
-        for a in sfp.getChildren():
+        for a in sfp:
             thisTime = a.getAge() / units.DAYS_PER_YEAR + a.p.chargeTime
 
             if thisTime != lastTime:

armi/physics/fuelCycle/fuelHandlers.py

@@ -654,7 +654,7 @@ def _getAssembliesInRings(
                 f"or otherwise instantiate a SpentFuelPool object as r.excore['sfp']"
             )
         else:
-            sfpAssems = self.r.excore["sfp"].getChildren()
+            sfpAssems = list(self.r.excore["sfp"])
 
         assemblyList = [[] for _i in range(len(ringList))]  # empty lists for each ring
         if exclusions is None:

armi/physics/fuelCycle/tests/test_fuelHandlers.py

@@ -468,7 +468,7 @@ def runShuffling(self, fh):
             self.r.p.cycle = cycle
             fh.cycle = cycle
             fh.manageFuel(cycle)
-            for a in self.r.excore["sfp"].getChildren():
+            for a in self.r.excore["sfp"]:
                 self.assertEqual(a.getLocation(), "SFP")
             for b in self.r.core.getBlocks(Flags.FUEL):
                 self.assertGreater(b.p.kgHM, 0.0, "b.p.kgHM not populated!")
@@ -498,7 +498,7 @@ def test_repeatShuffles(self):
         ensure repeatability.
         """
         # check labels before shuffling:
-        for a in self.r.excore["sfp"].getChildren():
+        for a in self.r.excore["sfp"]:
             self.assertEqual(a.getLocation(), "SFP")
 
         # do some shuffles
@@ -532,7 +532,7 @@ def test_repeatShuffles(self):
         # make sure the shuffle was repeated perfectly.
         for a in self.r.core.getAssemblies():
             self.assertEqual(a.getName(), firstPassResults[a.getLocation()])
-        for a in self.r.excore["sfp"].getChildren():
+        for a in self.r.excore["sfp"]:
             self.assertEqual(a.getLocation(), "SFP")
 
         # Do some cleanup, since the fuelHandler Interface has code that gets
@@ -811,7 +811,7 @@ def test_dischargeSwap(self):
 
         # grab an arbitrary fuel assembly from the core and from the SFP
         a1 = self.r.core.getAssemblies(Flags.FUEL)[0]
-        a2 = self.r.excore["sfp"].getChildren(Flags.FUEL)[0]
+        a2 = self.r.excore["sfp"].getChildrenWithFlags(Flags.FUEL)[0]
 
         # grab the stationary blocks pre swap
         a1PreSwapStationaryBlocks = [

armi/physics/neutronics/globalFlux/globalFluxInterface.py

@@ -277,7 +277,7 @@ def getTightCouplingValue(self):
             return self.r.core.p.keff
         if self.coupler.parameter == "power":
             scaledCorePowerDistribution = []
-            for a in self.r.core.getChildren():
+            for a in self.r.core:
                 scaledPower = []
                 assemPower = sum(b.p.power for b in a)
                 for b in a:

armi/physics/neutronics/globalFlux/tests/test_globalFluxInterface.py

@@ -286,7 +286,7 @@ def test_getTightCouplingValue(self):
         self._setTightCouplingTrue()
         self.assertEqual(self.gfi.getTightCouplingValue(), 1.0)  # set in setUp
         self.gfi.coupler.parameter = "power"
-        for a in self.r.core.getChildren():
+        for a in self.r.core:
             for b in a:
                 b.p.power = 10.0
         self.assertEqual(

armi/physics/neutronics/isotopicDepletion/crossSectionTable.py

@@ -223,7 +223,7 @@ def makeReactionRateTable(obj, nuclides: List = None):
         for nucName in nuclides
     }
 
-    for armiObject in obj.getChildren():
+    for armiObject in obj:
         for nucName in nuclides:
             rxnRates = armiObject.getReactionRates(nucName, nDensity=1.0)
             for rxName, rxRate in rxnRates.items():

armi/reactor/assemblies.py

@@ -326,10 +326,8 @@ def getPinPlenumVolumeInCubicMeters(self):
         -------
         This is a bit design-specific for pinned assemblies
         """
-        plenumBlocks = self.getBlocks(Flags.PLENUM)
-
         plenumVolume = 0.0
-        for b in plenumBlocks:
+        for b in self.iterChildrenWithFlags(Flags.PLENUM):
             cladId = b.getComponent(Flags.CLAD).getDimension("id")
             length = b.getHeight()
             plenumVolume += (
@@ -339,7 +337,7 @@ def getPinPlenumVolumeInCubicMeters(self):
 
     def getAveragePlenumTemperature(self):
         """Return the average of the plenum block outlet temperatures."""
-        plenumBlocks = self.getBlocks(Flags.PLENUM)
+        plenumBlocks = self.iterChildrenWithFlags(Flags.PLENUM)
         plenumTemps = [b.p.THcoolantOutletT for b in plenumBlocks]
 
         # no plenum blocks, use the top block of the assembly for plenum temperature
@@ -814,7 +812,7 @@ def dump(self, fName=None):
         with open(fName, "w") as pkl:
             pickle.dump(self, pkl)
 
-    def getBlocks(self, typeSpec=None, exact=False):
+    def getBlocks(self, typeSpec=None, exact=False) -> list[blocks.Block]:
         """
         Get blocks in an assembly from bottom to top.
 
@@ -831,9 +829,10 @@ def getBlocks(self, typeSpec=None, exact=False):
             List of blocks.
         """
         if typeSpec is None:
-            return self.getChildren()
+            items = iter(self)
         else:
-            return self.getChildrenWithFlags(typeSpec, exactMatch=exact)
+            items = self.iterChildrenWithFlags(typeSpec, exact)
+        return list(items)
 
     def getBlocksAndZ(self, typeSpec=None, returnBottomZ=False, returnTopZ=False):
         """
@@ -881,26 +880,40 @@ def getBlocksAndZ(self, typeSpec=None, returnBottomZ=False, returnTopZ=False):
         return zip(blocks, zCoords)
 
     def hasContinuousCoolantChannel(self):
-        return all(
-            b.containsAtLeastOneChildWithFlags(Flags.COOLANT) for b in self.getBlocks()
-        )
+        return all(b.containsAtLeastOneChildWithFlags(Flags.COOLANT) for b in self)
 
     def getFirstBlock(self, typeSpec=None, exact=False):
-        bs = self.getBlocks(typeSpec, exact=exact)
-        if bs:
-            return bs[0]
+        """Find the first block that matches the spec.
+
+        Parameters
+        ----------
+        typeSpec : flag or list of flags, optional
+            Specification to require on the returned block.
+        exact : bool, optional
+            Require block to exactly match ``typeSpec``
+
+        Returns
+        -------
+        Block or None
+            First block that matches if such a block could be found.
+        """
+        if typeSpec is None:
+            items = iter(self)
         else:
+            items = self.iterChildrenWithFlags(typeSpec, exact)
+        try:
+            # Create an iterator and attempt to advance it to the first value.
+            return next(items)
+        except StopIteration:
+            # No items found in the iteration -> no blocks match the request
             return None
 
     def getFirstBlockByType(self, typeName):
-        bs = [
-            b
-            for b in self.getChildren(deep=False)
-            if isinstance(b, blocks.Block) and b.getType() == typeName
-        ]
-        if bs:
-            return bs[0]
-        return None
+        blocks = filter(lambda b: b.getType() == typeName, self)
+        try:
+            return next(blocks)
+        except StopIteration:
+            return None
 
     def getBlockAtElevation(self, elevation):
         """
@@ -1191,7 +1204,11 @@ def countBlocksWithFlags(self, blockTypeSpec=None):
         blockCounter : int
             number of blocks of this type
         """
-        return len(self.getBlocks(blockTypeSpec))
+        if blockTypeSpec is None:
+            items = iter(self)
+        else:
+            items = self.iterChildrenWithFlags(blockTypeSpec)
+        return sum(1 for _ in items)
 
     def getDim(self, typeSpec, dimName):
         """

armi/reactor/blocks.py

@@ -627,7 +627,7 @@ def getArea(self, cold=False):
             return area
 
         a = 0.0
-        for c in self.getChildren():
+        for c in self:
             myArea = c.getArea(cold=cold)
             a += myArea
         fullArea = a
@@ -930,7 +930,7 @@ def add(self, c):
         self._updatePitchComponent(c)
 
     def removeAll(self, recomputeAreaFractions=True):
-        for c in self.getChildren():
+        for c in list(self):
             self.remove(c, recomputeAreaFractions=False)
         if recomputeAreaFractions:  # only do this once
             self.getVolumeFractions()
@@ -1283,7 +1283,7 @@ def getPinPitch(self, cold=False):
     def getDimensions(self, dimension):
         """Return dimensional values of the specified dimension."""
         dimVals = set()
-        for c in self.getChildren():
+        for c in self:
             try:
                 dimVal = c.getDimension(dimension)
             except parameters.ParameterError:
@@ -1532,7 +1532,7 @@ def getPinLocations(self) -> list[grids.IndexLocation]:
         :meth:`getPinCoordinates` - companion for this method.
         """
         items = []
-        for clad in self.getChildrenWithFlags(Flags.CLAD):
+        for clad in self.iterChildrenWithFlags(Flags.CLAD):
             if isinstance(clad.spatialLocator, grids.MultiIndexLocation):
                 items.extend(clad.spatialLocator)
             else:

armi/reactor/blueprints/__init__.py

@@ -519,7 +519,7 @@ def _checkAssemblyAreaConsistency(self, cs):
                     runLog.error("CURRENT COMPARISON BLOCK:")
                     b.printContents(includeNuclides=False)
 
-                    for c in b.getChildren():
+                    for c in b:
                         runLog.error(
                             "{0} area {1} effective area {2}"
                             "".format(c, c.getArea(), c.getVolume() / b.getHeight())

armi/reactor/components/__init__.py

@@ -383,7 +383,7 @@ def _deriveVolumeAndArea(self):
         # Determine the volume/areas of the non-derived shape components within the parent.
         siblingVolume = 0.0
         siblingArea = 0.0
-        for sibling in self.parent.getChildren():
+        for sibling in self.parent:
             if sibling is self:
                 continue
             elif not self and isinstance(sibling, DerivedShape):

armi/reactor/components/component.py

@@ -964,7 +964,7 @@ def clearLinkedCache(self):
     def getLinkedComponents(self):
         """Find other components that are linked to this component."""
         dependents = []
-        for child in self.parent.getChildren():
+        for child in self.parent:
             for dimName in child.DIMENSION_NAMES:
                 isLinked = child.dimensionIsLinked(dimName)
                 if isLinked and child.p[dimName].getLinkedComponent() is self: