From 38550f3bf87f6400ea577fcfab51545911e21e08 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Wed, 29 Nov 2017 06:34:15 -0600
Subject: [PATCH] Use sorted() to sort XML elements in geometry.xml

---
 openmc/clean_xml.py | 67 ---------------------------------------------
 openmc/geometry.py  |  7 +++--
 2 files changed, 5 insertions(+), 69 deletions(-)

diff --git a/openmc/clean_xml.py b/openmc/clean_xml.py
index 2497d0d9f..d1002070b 100644
--- a/openmc/clean_xml.py
+++ b/openmc/clean_xml.py
@@ -1,70 +1,3 @@
-def sort_xml_elements(tree):
-
-    # Retrieve all children of the root XML node in the tree
-    elements = list(tree)
-
-    # Initialize empty lists for the sorted and comment elements
-    sorted_elements = []
-
-    # Initialize an empty set of tags (e.g., Surface, Cell, and Lattice)
-    tags = set()
-
-    # Find the unique tags in the tree
-    for element in elements:
-        tags.add(element.tag)
-
-    # Initialize an empty list for the comment elements
-    comment_elements = []
-
-    # Find the comment elements and record their ordering within the
-    # tree using a precedence with respect to the subsequent nodes
-    for index, element in enumerate(elements):
-        next_element = None
-
-        if 'Comment' in str(element.tag):
-
-            if index < len(elements)-1:
-                next_element = elements[index+1]
-
-            comment_elements.append((element, next_element))
-
-    # Now iterate over all tags and order the elements within each tag
-    for tag in sorted(list(tags)):
-
-        # Retrieve all of the elements for this tag
-        try:
-            tagged_elements = tree.findall(tag)
-        except:
-            continue
-
-        # Initialize an empty list of tuples to sort (id, element)
-        tagged_data = []
-
-        # Retrieve the IDs for each of the elements
-        for element in tagged_elements:
-            key = element.get('id')
-
-            # If this element has an "ID" tag, append it to the list to sort
-            if key is not None:
-                tagged_data.append((int(key), element))
-
-        # Sort the elements according to the IDs for this tag
-        tagged_data.sort()
-        sorted_elements.extend(list(item[-1] for item in tagged_data))
-
-    # Add the comment elements while preserving the original precedence
-    for element, next_element in comment_elements:
-        index = sorted_elements.index(next_element)
-        sorted_elements.insert(index, element)
-
-    # Remove all of the sorted elements from the tree
-    for element in sorted_elements:
-        tree.remove(element)
-
-    # Add the sorted elements back to the tree in the proper order
-    tree.extend(sorted_elements)
-
-
 def clean_xml_indentation(element, level=0, spaces_per_level=2):
     """
     copy and paste from http://effbot.org/zone/elementent-lib.htm#prettyprint
diff --git a/openmc/geometry.py b/openmc/geometry.py
index 88c9f1834..0738ae615 100644
--- a/openmc/geometry.py
+++ b/openmc/geometry.py
@@ -5,7 +5,7 @@ from xml.etree import ElementTree as ET
 from six import string_types
 
 import openmc
-from openmc.clean_xml import sort_xml_elements, clean_xml_indentation
+from openmc.clean_xml import clean_xml_indentation
 from openmc.checkvalue import check_type
 
 
@@ -81,8 +81,11 @@ class Geometry(object):
         root_element = ET.Element("geometry")
         self.root_universe.create_xml_subelement(root_element)
 
+        # Sort the elements in the file
+        root_element[:] = sorted(root_element, key=lambda x: (
+            x.tag, int(x.get('id'))))
+
         # Clean the indentation in the file to be user-readable
-        sort_xml_elements(root_element)
         clean_xml_indentation(root_element)
 
         # Write the XML Tree to the geometry.xml file