[WIP] Parse Separations and Mixer archetypes properly

tests/test_commodity_dictionary.py

@@ -24,6 +24,31 @@ def test_commod_names_outcomodity():
                                                       ":cycamore:Reactor")
 
 
+@pytest.mark.parametrize("l,uitype,path", [("outcommodity", "incommodity", 
+                                            None),
+                                           (['oneormore', ['pair', ['pair',
+                                             'double', 'double'],['oneormore',
+                                             'incommodity','double']]],
+                                             'incommodity',[1, 2, 1])
+                                       ])
+def test_search_var_recursive(l, uitype, path):
+
+    assert path == cd.search_var_recursive(l, uitype)
+
+
+@pytest.mark.parametrize("data,p,alias", [(['in_streams', ['stream', ['info',
+                                           'mixing_ratio', 'buf_size'],
+                                           [['commodities', 'item'], 
+                                           'commodity', 'pref']]], [1, 2, 1],
+                                           'commodity'),
+                                          (['a', [['b', 'c'], 'd']], [1, 0,
+                                          0], 'b')
+                                         ])
+def test_find_alias(data, p, alias):
+
+    assert alias == cd.find_alias(data, p)
+
+
 def test_build_facility_dictionary():
     exp = {':agents:KFacility': (['in_commod'], ['out_commod']),
            ':agents:NullInst': ([], []),
@@ -39,12 +64,12 @@ def test_build_facility_dictionary():
                                   'topup_commod'], ['outcommod']),
            ':cycamore:GrowthRegion': ([], []),
            ':cycamore:ManagerInst': ([], []),
-           ':cycamore:Mixer': ([], ['out_commod']),
+           ':cycamore:Mixer': (['commodity'], ['out_commod']),
            ':cycamore:Reactor': (['fuel_incommods', 'pref_change_commods',
                                   'recipe_change_commods'],
                                  ['fuel_outcommods']),
            ':cycamore:Separations': (['feed_commods'], ['leftover_commod',
-                                      'streams_']),
+                                      'commod']),
            ':cycamore:Sink': (['in_commods'], []),
            ':cycamore:Source': ([], ['outcommod']),
            ':cycamore:Storage': (['in_commods'], ['out_commods'])}

trailmap/commodity_dictionary.py

@@ -27,7 +27,7 @@ def get_commod_names(metadata, uitype, agent):
     '''Return all archetypes and their aliases for a given uitype.
 
     inputs:
-    - metadata
+    - metadata: Cyclus metadata
     - uitype: a string. Example, "incommodity"
     - agent: a string. Example, ":cycamore:Enrichment"
 
@@ -40,14 +40,77 @@ def get_commod_names(metadata, uitype, agent):
 
     for var in agent_data:
         for param in agent_data[var]:
-            if param == "uitype" and uitype in agent_data[var][param]:
-                aliases.append(var)
+            if param == "uitype":
+                if uitype in agent_data[var][param]: #typical for archetypes
+                    aliases.append(var)
+                #if streams are present
+                path = search_var_recursive(agent_data[var][param], uitype)
+                if path is not None:
+                    aliases.append(find_alias(agent_data[var]["alias"], path))
+
+    #drop the "val" (signifing an interleave) and "streams_" aliases 
+    try:
+        while True:
+            aliases.remove("val")
+    except:
+        pass
+    try:
+        while True:
+            aliases.remove('streams_')
+    except:
+        pass
 
     return aliases
 
 
+def search_var_recursive(var, uitype):
+    '''Finds the path within streams_ or similar tree that matches desired
+    uitype. Searches recursively
+
+    inputs:
+        - var: a nested list from an archetype's metadata
+        - uitype: a string. Example, "incommodity"
+
+    outputs:
+        - path: the path within var that locates uitype
+    '''
+    for index,item in enumerate(var):
+        if item == uitype:
+            return [index]
+        if isinstance(item, list):
+            path = search_var_recursive(item, uitype)
+            if path:
+                return [index] + path
+
+
+def find_alias(var, path):
+    '''Given path to locate uitype, searches archetype aliases to locate
+    matching alias
+
+    inputs:
+        - var: a nested list of aliases
+        - path: the path to search for the desired alias
+
+    outputs:
+        - alias: a string
+    '''
+    if len(path) == 1:
+        return var[path[0]]
+    else:
+        return find_alias(var[path[0]], path[1:])
+
+
 def build_facility_dictionary(metadata, archetypes):
-    '''Identify commodities for each available archetype'''
+    '''Identify commodities for each available archetype
+
+    inputs:
+        - metadata: Cyclus metadata
+        - archetypes: a list of archetypes to use
+
+    outputs:
+        - archetype_commods: a dictionary with the Cyclus archetypes available
+        and the names of their incommodities and outcommodities 
+    '''
     archetype_commods = {}
 
     for archetype in archetypes:
@@ -56,6 +119,8 @@ def build_facility_dictionary(metadata, archetypes):
         outcommods = get_commod_names(metadata["annotations"], "outcommodity",
                                       archetype)
 
+
+
         archetype_commods.update({archetype: (incommods,
                                               outcommods)})
 

trailmap/parse_input.py

@@ -49,19 +49,27 @@ def get_facility_and_commod_names(root, input_archetypes,
         facility_archetype = facility.find('config/').tag
         facility_module = input_archetypes[facility_archetype]
 
-        (in_tags, out_tags) = commodity_dictionary[facility_module]
-
-        facility_in_commods = []
-        facility_out_commods = []
+        if facility_module == ':cycamore:Separations':
+            (facility_in_commods,
+             facility_out_commods) = find_sep_commod(facility,
+                                                     facility_archetype)
+        elif facility_module == ':cycamore:Mixer':
+            (facility_in_commods,
+             facility_out_commods) = find_mixer_commod(facility,
+                                                       facility_archetype)
+        else:
+            (in_tags, out_tags) = commodity_dictionary[facility_module]
+            facility_in_commods = []
+            facility_out_commods = []
 
-        for archetype_var in facility.find('.config/' + facility_archetype):
-            in_commods = find_commod(archetype_var, in_tags)
-            if in_commods is not None:
-                facility_in_commods.extend(in_commods)
+            for archetype_var in facility.find('.config/'+facility_archetype):
+                in_commods = find_commod(archetype_var, in_tags)
+                if in_commods is not None:
+                    facility_in_commods.extend(in_commods)
 
-            out_commods = find_commod(archetype_var, out_tags)
-            if out_commods is not None:
-                facility_out_commods.extend(out_commods)
+                out_commods = find_commod(archetype_var, out_tags)
+                if out_commods is not None:
+                    facility_out_commods.extend(out_commods)
 
         facility_dict_in[facility_name] = facility_in_commods
         facility_dict_out[facility_name] = facility_out_commods
@@ -112,3 +120,46 @@ def find_commod(archetype_tag, commod_tags):
         return commod_list
 
     return commod_list
+
+
+def find_sep_commod(facility, facility_archetype):
+    '''Searches for commodities within a Separations facility, which uses a
+    different xml schema than other Cyclus archetypes
+    '''
+    in_tags = ['feed_commods']
+    leftover_tags = ['leftover_commod']
+
+    facility_in_commods = []
+    facility_out_commods = []
+
+    for archetype_var in facility.find('.config/' + facility_archetype):
+        in_commods = find_commod(archetype_var, in_tags)
+        if in_commods is not None:
+            facility_in_commods.extend(in_commods)
+        if archetype_var.tag == 'streams':
+            for commod in archetype_var.findall('./item/commod'):
+                facility_out_commods.append(commod.text)
+        out_commods = find_commod(archetype_var, leftover_tags)
+        if out_commods is not None:
+            facility_out_commods.extend(out_commods)
+
+    return facility_in_commods, facility_out_commods
+
+
+def find_mixer_commod(facility, facility_archetype):
+    '''Searches for commodities within a Mixer facility, which uses a
+    different xml schema than other Cyclus archetypes
+    '''
+    out_tags = ['out_commod']
+    facility_in_commods = []
+    facility_out_commods = []
+
+    for archetype_var in facility.find('.config/' + facility_archetype):
+        if archetype_var.tag == 'in_streams':
+            for commod in archetype_var.findall('./stream/commodities/item/commodity'):
+                facility_in_commods.append(commod.text)
+        out_commods = find_commod(archetype_var, out_tags)
+        if out_commods is not None:
+            facility_out_commods.extend(out_commods)
+
+    return facility_in_commods, facility_out_commods