" Added company file capabilities "

wetb/hawc2/hawc2_simulation.py

@@ -49,6 +49,10 @@ class Hawc2_Simulation(object):
         # The directories for the source and the destination
         self.source_directory = None
         self.write_directory = None
+        
+        #CMOS: company file.
+        self.company_files = None
+
         # indicates that the sensors in the source htc file are normalized
         self.source_sensors_normalized = False
         # indicates that the object should keep the sensors normalized
@@ -149,6 +153,22 @@ class Hawc2_Simulation(object):
         if self.write_directory == self.source_directory and not force_write:
             raise Exception('Cannot write the contents into the source directory without the write being forced')
 
+        #CMOS: Copying company files:
+        if not self.company_files is None:
+            company_index = 0
+            for company_file in self.company_files:
+                if self.company_destinations[company_index] is None:
+                    shutil.copy(company_file,self.write_directory)
+                else:
+                    try:
+                        shutil.copy(company_file,self.write_directory+self.company_destinations[company_index])
+                    except:
+                        print('CMOS: Something went wrong when copying company files. The file is put in the writing directory')
+                        shutil.copy(company_file,self.write_directory)
+                company_index += 1
+        #CMOS - end
+
+
         # Get the directory
         old_home = os.getcwd()
 
@@ -205,6 +225,18 @@ class Hawc2_Simulation(object):
     # This will set the write directory
     def set_write_directory(self, write_directory):
         self.write_directory = os.path.abspath(os.path.realpath(write_directory))
+    
+    #CMOS: Set company file:
+    def set_company_file(self, company_file, destination=None):
+        if self.company_files is None:
+            self.company_files = list()
+            self.company_files.append(os.path.abspath(os.path.realpath(company_file)))
+            self.company_destinations = list()
+            self.company_destinations.append(destination)
+        else:
+            self.company_files.append(os.path.abspath(os.path.realpath(company_file)))
+            self.company_destiantions.append(destination)
+    #CMOS - end
 
     # Tests whether the input for the current state has already been written
     def does_input_exist(self):