From 5f245f2e7da760641f766a3dddccfbe6aac931d3 Mon Sep 17 00:00:00 2001
From: Ruben1729 <ruben.e.sanchez.a@gmail.com>
Date: Wed, 14 Feb 2024 17:06:21 -0500
Subject: [PATCH] Cleanup for matsim export factory

---
 matsim.py | 54 +++++++++++++++++++++++++-----------------------------
 1 file changed, 25 insertions(+), 29 deletions(-)

diff --git a/matsim.py b/matsim.py
index 2a6ac3b..5dfc26a 100644
--- a/matsim.py
+++ b/matsim.py
@@ -23,7 +23,7 @@ class Matsim:
       'facility': []
     }
 
-  def _export(self):
+  def export(self):
     self._export_facilities()
     self._export_network()
     self._export_population()
@@ -96,19 +96,10 @@ class Matsim:
     )
     gdf.to_file("input_files/buildings_shapefile.shp")
 
-    # Convert the Python dictionary to an XML string
+    # Write xml content to file
     xml_content = etree.tostring(facilities_xml, pretty_print=True, encoding='UTF-8').decode('utf-8')
-
-    # Write the XML to the file
     output_file = f"{self._output_file_path}/{self._city.name}_facilities.xml"
-    with open(output_file, 'w') as file:
-      file.write("<?xml version=\"1.0\" encoding=\"utf-8\"?>\n")
-      file.write(f"<!DOCTYPE facilities SYSTEM \"{FACILITIES_DTD}\">\n")
-      file.write(xml_content)
-
-    with open(output_file, 'rb') as f_in:
-      with gzip.open(output_file + '.gz', 'wb') as f_out:
-        shutil.copyfileobj(f_in, f_out)
+    _save_xml(output_file, xml_content, f"<!DOCTYPE facilities SYSTEM \"{FACILITIES_DTD}\">\n", True)
 
   def _export_network(self):
     java_path = "java"
@@ -194,19 +185,10 @@ class Matsim:
 
           id += 1
 
-    # Convert the Python dictionary to an XML string
+    # Write xml content to file
     xml_content = etree.tostring(population, pretty_print=True, encoding='UTF-8').decode('utf-8')
-
-    # Write the XML to the file
     output_file = f"{self._output_file_path}/{self._city.name}_population.xml"
-    with open(output_file, 'w') as file:
-      file.write("<?xml version=\"1.0\" encoding=\"utf-8\"?>\n")
-      file.write(f"<!DOCTYPE population SYSTEM \"{POPULATION_DTD}\">\n")
-      file.write(xml_content)
-
-    with open(output_file, 'rb') as f_in:
-      with gzip.open(output_file + '.gz', 'wb') as f_out:
-        shutil.copyfileobj(f_in, f_out)
+    _save_xml(output_file, xml_content, f"<!DOCTYPE population SYSTEM \"{POPULATION_DTD}\">\n", True)
 
   def _export_config(self):
     root = etree.Element('config')
@@ -377,14 +359,10 @@ class Matsim:
     # Defines all the modes available, including chain-based modes, seperated by commas
     _add_param(subtour_module, 'modes', 'car,walk')
 
+    # Write xml content to file
     xml_content = etree.tostring(root, pretty_print=True, encoding='UTF-8').decode('utf-8')
-
-    # Write the XML to the file
     output_file = f"{self._output_file_path}/{self._city.name}_config.xml"
-    with open(output_file, 'w') as file:
-      file.write("<?xml version=\"1.0\" encoding=\"utf-8\"?>\n")
-      file.write(f"<!DOCTYPE config SYSTEM \"{CONFIG_DTD}\">\n")
-      file.write(xml_content)
+    _save_xml(output_file, xml_content, f"<!DOCTYPE config SYSTEM \"{CONFIG_DTD}\">\n")
 
 
 def _add_param(parent, name, value):
@@ -427,3 +405,21 @@ def _convert_schedules(building_schedules):
       })
 
   return converted_schedules
+
+
+def _save_xml(output_file, xml_content, xml_dtd=None, zipped=None):
+  if zipped is None:
+    zipped = False
+
+  if xml_dtd is None:
+    xml_dtd = ''
+
+  with open(output_file, 'w') as file:
+    file.write("<?xml version=\"1.0\" encoding=\"utf-8\"?>\n")
+    file.write(xml_dtd)
+    file.write(xml_content)
+
+  if zipped:
+    with open(output_file, 'rb') as f_in:
+      with gzip.open(output_file + '.gz', 'wb') as f_out:
+        shutil.copyfileobj(f_in, f_out)