Make Process.default_state recursive

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## v1.5.8
+* (#227) Make `Process.default_state` parse ports schema recursively.
+
 ## v1.5.7
 * (#226) Drop Python 3.8 in favor of 3.11 and update dependencies.
 

setup.py

@@ -2,7 +2,7 @@
 from setuptools import setup
 
 
-VERSION = '1.5.7'
+VERSION = '1.5.8'
 
 
 if __name__ == '__main__':

vivarium/composites/toys.py

@@ -614,6 +614,84 @@ def generate_topology(
                 'port2': ('store_C',)}}
 
 
+class MixedLevelPorts(Process):
+    """Process with ports that configure top-level, second-level, third-level,
+    and fourth-level stores. Used to test recursive Process.default_state"""
+    name = 'mixed_level_ports'
+
+    def ports_schema(self) -> Schema:
+        return {
+            'a': {'_default': 0},
+            'b': {
+                'c': {'_default': 1},
+                'd': {'_default': 2}},
+            'e': {
+                'f': {
+                    'g': {'_default': 3},
+                    'h': {'_default': 4},
+                    'i': {'_default': 5}}},
+            'j': {
+                'k': {
+                    'l': {
+                        'm': {'_default': 6},
+                        'n': {'_default': 7},
+                        'o': {'_default': 8}},
+                    'p': {'_default': 9},
+                    'q': {'r': {'_default': 10}}}},
+        }
+
+    def next_update(
+            self, timestep: Union[float, int], states: State) -> Update:
+        return {}
+
+
+class MixedLevelPortsComposer(Composer):
+
+    def generate_processes(
+            self,
+            config: Optional[dict]
+    ) -> Dict[str, MixedLevelPorts]:
+        return {'mixed_level_ports': MixedLevelPorts()}
+
+    def generate_topology(
+            self,
+            config: Optional[dict] = None
+    ) -> Topology:
+        return {
+            'mixed_level_ports': {
+                'a': ('store_a',),
+                'b': ('store_b',),
+                'e': ('store_e',),
+                'j': ('store_j',)}}
+
+
+def test_default_state() -> None:
+    composer = MixedLevelPortsComposer()
+    composite = composer.generate()
+    default_state = composite.default_state()
+
+    expected_default_state = {
+        'store_a': 0,
+        'store_b': {
+            'c': 1,
+            'd': 2},
+        'store_e': {
+            'f': {
+                'g': 3,
+                'h': 4,
+                'i': 5}},
+        'store_j': {
+            'k': {
+                'l': {
+                    'm': 6,
+                    'n': 7,
+                    'o': 8},
+                'p': 9,
+                'q': {'r': 10}}},
+    }
+    assert default_state == expected_default_state
+
+
 def test_override() -> None:
     config = {
         '_schema': {
@@ -928,6 +1006,7 @@ def generate_topology(self, config):
 
 
 if __name__ == '__main__':
+    test_default_state() # pragma: no cover
     test_composite_initial_state()  # pragma: no cover
     test_composite_parameters()  # pragma: no cover
     test_composite_merge()  # pragma: no cover

vivarium/core/process.py

@@ -464,7 +464,7 @@ def calculate_timestep(self, states: Optional[State]) -> Union[float, int]:
         return self.parameters['timestep']
 
     def default_state(self) -> State:
-        """Get the default values of the variables in each port.
+        """Recursively get the default values of the variables in each port.
 
         The default values are computed based on the schema.
 
@@ -473,14 +473,22 @@ def default_state(self) -> State:
             value to that variable.
         """
         schema = self.get_schema()
-        state: State = {}
-        for port, states in schema.items():
-            for key, value in states.items():
-                if '_default' in value:
-                    if port not in state:
-                        state[port] = {}
-                    state[port][key] = value['_default']
-        return state
+
+        def get_defaults(d: dict) -> Optional[State]:
+            defaults = {}
+            for k, v in d.items():
+                if isinstance(v, dict):
+                    def_val = v.get('_default', get_defaults(v))
+                    if def_val is not None:
+                        defaults[k] = def_val
+            if len(defaults) > 0:
+                return defaults
+            return None
+
+        default_state = get_defaults(schema)
+        if default_state is None:
+            return {}
+        return default_state
 
     def is_deriver(self) -> bool:
         """Check whether this process is a :term:`deriver`.