update helper functions

Author: snowman2

spt_compute/imports/helper_functions.py

@@ -1,32 +1,36 @@
 # -*- coding: utf-8 -*-
-##
-##  helper_functions.py
-##  spt_compute
-##
-##  Created by Alan D. Snow.
-##  Copyright © 2015-2016 Alan D Snow. All rights reserved.
-##  License: BSD-3 Clause
+#
+#  helper_functions.py
+#  spt_ecmwf_autorapid_process
+#
+#  Created by Alan D. Snow
+#  License: BSD-3 Clause
 
 import datetime
 from glob import glob
-from netCDF4 import Dataset, num2date  # http://unidata.github.io/netcdf4-python/
 import os
 import re
 from shutil import rmtree
 import sys
 
 
-#----------------------------------------------------------------------------------------
+# ----------------------------------------------------------------------------------------
 # HELPER FUNCTIONS
-#----------------------------------------------------------------------------------------
+# ----------------------------------------------------------------------------------------
 class CaptureStdOutToLog(object):
-    def __init__(self, out_file_path):
-        self.out_file_path = out_file_path
+    def __init__(self, log_file_path, error_file_path=None):
+        self.log_file_path = log_file_path
+        self.error_file_path = error_file_path
+        if error_file_path is None:
+            self.error_file_path = "{0}.err".format(os.path.splitext(log_file_path)[0])
+
     def __enter__(self):
         self._stdout = sys.stdout
         self._stderr = sys.stderr
-        sys.stdout = sys.stderr = open(self.out_file_path, 'w')
+        sys.stdout = open(self.log_file_path, 'w')
+        sys.stderr = open(self.error_file_path, 'w')
         return self
+
     def __exit__(self, *args):
         sys.stdout.close()
         sys.stdout = self._stdout
@@ -46,43 +50,44 @@ def case_insensitive_file_search(directory, pattern):
         raise
 
 
-def clean_main_logs(main_log_directory, prepend="spt_compute_ecmwf_",
-                    lock_file_name="spt_compute_ecmwf_run_info_lock.txt",
-                    log_file_path=""):
+def clean_main_logs(main_log_directory, prepend="rapid_", log_file_path=""):
     """
-    Removes old log files older than one week old in main log directory
+    This removes main logs older than three days old
     """
     date_today = datetime.datetime.utcnow()
-    week_timedelta = datetime.timedelta(7)
+    week_timedelta = datetime.timedelta(3)
+
+    # clean up log files
     main_log_files = [f for f in os.listdir(main_log_directory) if
                       not os.path.isdir(os.path.join(main_log_directory, f))
                       and not log_file_path.endswith(f)
-                      and f != lock_file_name]
+                      and (f.endswith('log') or f.endswith('err'))]
 
     for main_log_file in main_log_files:
         try:
-            log_datetime = datetime.datetime.strptime(main_log_file, "{0}%y%m%d%H%M%S.log".format(prepend))
-            if (date_today-log_datetime > week_timedelta):
+            log_datetime = datetime.datetime.strptime(main_log_file[:18],
+                                                      "{0}%y%m%d%H%M%S".format(
+                                                          prepend))
+            if date_today - log_datetime > week_timedelta:
                 os.remove(os.path.join(main_log_directory, main_log_file))
         except Exception as ex:
             print(ex)
             pass
 
 
-def clean_logs(condor_log_directory, main_log_directory, prepend="spt_compute_ecmwf_", log_file_path=""):
+def clean_logs(condor_log_directory, main_log_directory, prepend="rapid_", log_file_path=""):
     """
-    This removed logs older than one week old
+    This removes all logs older than three days old
     """
     date_today = datetime.datetime.utcnow()
-    week_timedelta = datetime.timedelta(7)
-    #clean up condor logs
+    week_timedelta = datetime.timedelta(3)
+    # clean up condor logs
     condor_dirs = [d for d in os.listdir(condor_log_directory) if
                    os.path.isdir(os.path.join(condor_log_directory, d))]
-
     for condor_dir in condor_dirs:
         try:
             dir_datetime = datetime.datetime.strptime(condor_dir[:11], "%Y%m%d.%H")
-            if (date_today-dir_datetime > week_timedelta):
+            if date_today-dir_datetime > week_timedelta:
                 rmtree(os.path.join(condor_log_directory, condor_dir))
         except Exception as ex:
             print(ex)
@@ -96,10 +101,11 @@ def find_current_rapid_output(forecast_directory, watershed, subbasin):
     Finds the most current files output from RAPID
     """
     if os.path.exists(forecast_directory):
-        basin_files = glob(os.path.join(forecast_directory,"Qout_%s_%s_*.nc" % (watershed, subbasin)))
-        if len(basin_files) >0:
+        basin_files = glob(os.path.join(forecast_directory,
+                                        "Qout_{0}_{1}_*.nc".format(watershed, subbasin)))
+        if len(basin_files) > 0:
             return basin_files
-    #there are none found
+    # there are none found
     return None
 
 
@@ -110,7 +116,7 @@ def get_valid_watershed_list(input_directory):
     valid_input_directories = []
     for directory in os.listdir(input_directory):
         if os.path.isdir(os.path.join(input_directory, directory)) \
-            and len(directory.split("-")) == 2:
+                and len(directory.split("-")) == 2:
             valid_input_directories.append(directory)
         else:
             print("{0} incorrectly formatted. Skipping ...".format(directory))
@@ -121,8 +127,8 @@ def get_date_timestep_from_forecast_folder(forecast_folder):
     """
     Gets the datetimestep from forecast
     """
-    #OLD: Runoff.20151112.00.netcdf.tar.gz
-    #NEW: Runoff.20160209.0.exp69.Fgrid.netcdf.tar
+    # OLD: Runoff.20151112.00.netcdf.tar.gz
+    # NEW: Runoff.20160209.0.exp69.Fgrid.netcdf.tar
     forecast_split = os.path.basename(forecast_folder).split(".")
     forecast_date_timestep = ".".join(forecast_split[1:3])
     return re.sub("[^\d.]+", "", forecast_date_timestep)
@@ -142,16 +148,14 @@ def get_datetime_from_forecast_folder(forecast_folder):
     :param forecast_folder:
     :return:
     """
-    return get_datetime_from_date_timestep(
-        get_date_timestep_from_forecast_folder(forecast_folder))
-
+    return get_datetime_from_date_timestep(get_date_timestep_from_forecast_folder(forecast_folder))
 
 def get_ensemble_number_from_forecast(forecast_name):
     """
     Gets the datetimestep from forecast
     """
-    #OLD: 20151112.00.1.205.runoff.grib.runoff.netcdf
-    #NEW: 52.Runoff.nc
+    # OLD: 20151112.00.1.205.runoff.grib.runoff.netcdf
+    # NEW: 52.Runoff.nc
     forecast_split = os.path.basename(forecast_name).split(".")
     if forecast_name.endswith(".205.runoff.grib.runoff.netcdf"):
         ensemble_number = int(forecast_split[2])
@@ -183,4 +187,4 @@ def log(message, severity):
     if severity in print_me:
         print("{0} {1}".format(severity, message))
     else:
-        raise Exception(message)
+        raise Exception(message)