source: cows_wps/trunk/process_modules/extract_uk_station_data.py @ 6124

"""
extract_uk_station_data.py
===================

Process extract_uk_station_data that holds the ExtractUKStationData class.

"""

import os, stat, time, sys

from cows_wps.process_handler.fileset import FileSet, FLAG
import cows_wps.process_handler.ddp_process_support as ddp_process_support
from cows_wps.process_handler.context.process_status import STATUS

# Import local modules
sys.path.append("/home/badc/software/datasets/ukmo-midas/scripts/extract")
# MIDAS Station search code
import getStations
import midasSubsetter


class ExtractUKStationData(object):

    def __call__(self, context):
        
        startTime = time.time()
        jobId = os.path.basename(context.processDir)

        # Parse the inputs
        ci = context.inputs
        Username = ci['Username']
        StationIDs = ci.get("StationIDs", []) 
        Delimiter = ci["Delimiter"]
        Counties = ci.get('Counties', [])
        BBox = ci.get("BBox", None)
        DataTypes = ci.get("DataTypes", [])

        # Get required start/end times
        StartDateTime = ci['StartDateTime']
        StartDateTime = self._revertDateTimeToLongString(StartDateTime)
        EndDateTime = ci['EndDateTime']
        EndDateTime = self._revertDateTimeToLongString(EndDateTime)

        ObsTableName = ci['ObsTableName'] 

        context.setStatus(STATUS.STARTED, 'Job is now running', 0)

        # Always need a FileSet, even if empty
        self.fileSet = context.outputs['FileSet'] = FileSet()

        stations_file = "uk_stations.txt"
        sf_path = context.processDir + '/outputs/' + stations_file

        # Get station IDs if not provided
        if StationIDs == []:
            # Check we have either counties or bbox to search domain on
            if Counties == [] and BBox == None:
                raise Exception("Invalid arguments provided. Must provide either a list of station IDs, a list of counties or a valid geographical bounding box.")

            # Call code to get Weather Stations and write file
            st_getter = getStations.StationIDGetter(Counties, bbox=BBox, dataTypes=DataTypes,
                       startTime=StartDateTime, endTime=EndDateTime, outputFile=sf_path, noprint=1)
            StationIDs = st_getter.stList
            StationIDs.sort()

        # else write the file one per station id per line
        else:
            StationIDs.sort()
            fout = open(sf_path, "w")
            fout.write("\r\n".join([str(st_id) for st_id in StationIDs])) 
            fout.close()

        self._addFileToFileSet(sf_path, "Station IDs file used for extraction of station data.", FLAG.DATA)

        context.outputs['ProcessSpecificContent'] = {"StationIDs": " ".join(StationIDs)}

        # Pretend that took 10% of time
        context.setStatus(STATUS.COMPLETED, 'The End', 10)

        # Now extract the data itself
        data_file_root = "station_data"
        if Delimiter == "comma":
            data_file = data_file_root + ".csv"
        else:
            data_file = data_file_root + ".txt"

        df_path = context.processDir + '/outputs/' + data_file

        # Need temp dir for big file extractions
        process_tmp_dir = context.processDir + '/tmp'

        midasSubsetter.MIDASSubsetter([ObsTableName], df_path, startTime=StartDateTime, 
                       endTime=EndDateTime, src_ids=StationIDs, tempDir=process_tmp_dir)

        self._addFileToFileSet(df_path, "Station data file.", FLAG.DATA)

        context.setStatus(STATUS.COMPLETED, 'The End', 100)
        completionTime = time.time()
        ddp_process_support.updateJobDetailsAfterCompletion(context, startTime, completionTime)


    def _addFileToFileSet(self, path, info, type=FLAG.DATA):
        "Adds file to output file set."
        f_size = os.path.getsize(path)
        output_basename = os.path.basename(path)
        self.fileSet.contents.append(FileSet(type, output_basename, f_size, info))


    def _revertDateTimeToLongString(self, dt):
        """
        Turns a date/time into a long string as needed by midas code.
        """
        return str(dt).replace("-", "").replace(" ", "").replace("T", "").replace(":", "")

    def dryRun(self, context):

        # Not implemented for sync jobs
        pass