Adding methods for obtaining nodes from Demographics objects by name. Final part of (Fix #690).

Author: ckirkman-IDM

emod_api/demographics/Demographics.py

@@ -460,6 +460,14 @@ def _all_nodes(self) -> List[Node]:
         default_node = [self.default_node] if hasattr(self, 'default_node') else []
         return self.nodes + default_node
 
+    @property
+    def _all_node_names(self) -> List[int]:
+        return [node.name for node in self._all_nodes]
+
+    @property
+    def _all_nodes_by_name(self) -> Dict[int, Node]:
+        return {node.name: node for node in self._all_nodes}
+
     @property
     def _all_node_ids(self) -> List[int]:
         return [node.id for node in self._all_nodes]
@@ -470,7 +478,7 @@ def _all_nodes_by_id(self) -> Dict[int, Node]:
 
     def get_node_by_id(self, node_id: int) -> Node:
         """
-        Returns the Node objects requested by their node id.
+        Returns the Node object requested by its node id.
 
         Args:
             node_id: a node_id to use in retrieving the requested Node object. None or 0 for 'the default node'.
@@ -505,6 +513,44 @@ def get_nodes_by_id(self, node_ids: List[int]) -> Dict[int, Node]:
         requested_nodes = {node_id: node for node_id, node in self._all_nodes_by_id.items() if node_id in node_ids}
         return requested_nodes
 
+    def get_node_by_name(self, node_name: str) -> Node:
+        """
+        Returns the Node object requested by its node name.
+
+        Args:
+            node_name: a node_name to use in retrieving the requested Node object. None for 'the default node'.
+
+        Returns:
+            a Node object
+        """
+        return list(self.get_nodes_by_name(node_names=[node_name]).values())[0]
+
+    def get_nodes_by_name(self, node_names: List[str]) -> Dict[str, Node]:
+        """
+        Returns the Node objects requested by their node name.
+
+        Args:
+            node_names: a list of node names to use in retrieving Node objects. None for 'the default node'.
+
+        Returns:
+            a dict with name: node entries
+        """
+        # replace a None name (default node) request with the default node's name
+        if node_names is None:
+            node_names = [self.default_node.name]
+        if None in node_names:
+            node_names.remove(None)
+            node_names.append(self.default_node.name)
+
+        missing_node_names = [node_name for node_name in node_names if node_name not in self._all_node_names]
+        if len(missing_node_names) > 0:
+            msg = ', '.join([str(node_name) for node_name in missing_node_names])
+            raise self.UnknownNodeException(f"The following node name(s) were requested but do not exist in this demographics "
+                                            f"object:\n{msg}")
+        requested_nodes = {node_name: node for node_name, node in self._all_nodes_by_name.items()
+                           if node_name in node_names}
+        return requested_nodes
+
     def SetMigrationPattern(self, pattern: str = "rwd"):
         """
         Set migration pattern. Migration is enabled implicitly.

tests/test_demog.py

@@ -24,6 +24,57 @@ def setUp(self) -> None:
         print(f"\n{self._testMethodName} started...")
         self.out_folder = manifest.demo_folder
 
+    def test_get_node_by_name(self):
+        from emod_api.demographics.Demographics import DemographicsBase
+
+        mars = Node.Node(lat=0, lon=0, pop=100, name='Mars')
+        venus = Node.Node(lat=0, lon=0, pop=100, name='Venus')
+        earth = Node.Node(lat=0, lon=0, pop=100, name='Earth')
+        nodes = [mars, venus]
+        demographics = Demographics.Demographics(nodes=nodes, default_node=earth)
+
+        # a non default node
+        node = demographics.get_node_by_name(node_name=mars.name)
+        self.assertEqual(node, nodes[0])
+
+        # the default node, checked explicitly then implicitly
+        node = demographics.get_node_by_name(node_name=earth.name)
+        self.assertEqual(node, earth)
+
+        node = demographics.get_node_by_name(node_name=None)
+        self.assertEqual(node, earth)
+
+        # a node name that does not exist (yet, at least!)
+        self.assertRaises(DemographicsBase.UnknownNodeException, demographics.get_node_by_name, node_name='Planet X')
+
+    def test_get_nodes_by_name(self):
+        from emod_api.demographics.Demographics import DemographicsBase
+
+        mars = Node.Node(lat=0, lon=0, pop=100, name='Mars')
+        venus = Node.Node(lat=0, lon=0, pop=100, name='Venus')
+        earth = Node.Node(lat=0, lon=0, pop=100, name='Earth')
+        nodes = [mars, venus]
+        demographics = Demographics.Demographics(nodes=nodes, default_node=earth)
+
+        # just getting some nodes, also checking explicit default node request, too.
+        nodes = demographics.get_nodes_by_name(node_names=['Mars', 'Earth'])
+        expected = {earth.name: earth, mars.name: mars}
+        self.assertEqual(nodes, expected)
+
+        # verify that a node name of None will yield the default node
+        nodes = demographics.get_nodes_by_name(node_names=['Mars', None])
+        expected = {earth.name: earth, mars.name: mars}
+        self.assertEqual(nodes, expected)
+
+        nodes = demographics.get_nodes_by_name(node_names=None)
+        expected = {earth.name: earth}
+        self.assertEqual(nodes, expected)
+
+        # a node name that does not exist (yet, at least!)
+        self.assertRaises(DemographicsBase.UnknownNodeException, demographics.get_nodes_by_name,
+                          node_names=['Planet X', earth.name])
+
+
     def test_demo_basic_node(self):
         out_filename = os.path.join(self.out_folder, "demographics_basic_node.json")
         demog = Demographics.from_template_node()
@@ -2609,7 +2660,6 @@ def test_SetEquilibriumVitalDynamicsFromWorldBank_EH_06(self):
         self.from_csv_with_country("Least developed countries: UN classification", 1980)
     def test_SetEquilibriumVitalDynamicsFromWorldBank_EH_07(self):
         self.from_csv_with_country("Turks and Caicos Islands", 1980)
-        
     #endregion
 if __name__ == '__main__':
     unittest.main()