root/spongenet/trunk/spongenet/query.py

#!/usr/bin/env python

"""
Create CSV spreadsheets from sponge data NetCDF files.

Usage:

   > python query.py path/to/config/file
   > python query.py -t
   > python query.py --test

Test silent import:

>>> import expand
"""

__author__ = "Chris Calloway"
__email__ = "cbc@chriscalloway.org"
__copyright__ = "Copyright 2010 UNC-CH Department of Marine Science"
__license__ = "GPL2"

import warnings
warnings.simplefilter("ignore", DeprecationWarning)

import sys
import os
import glob
from pycdf import CDF, CDFError
import time
import csv
import doctest
import unittest
from StringIO import StringIO

USAGE = "\n".join(__doc__.splitlines()[3:8])
TEST_PATH = os.path.join("tests", "query")


def _test():
    """
    Run doctests as unittest suite.

    Test silent import:

    >>> from expand import _test
    """

    suite = []
    suite.append(doctest.DocTestSuite())
    suite = unittest.TestSuite(suite)
    unittest.TextTestRunner().run(suite)

    return


def config_path():
    """
    Return the configuration file path from the command line.

    Supply too few arguments on command line:

    >>> save_stdout = sys.stdout
    >>> temp_stdout = StringIO()
    >>> sys.stdout = temp_stdout
    >>> sys.argv = []
    >>> _config_path = config_path()
    >>> sys.stdout = save_stdout
    >>> USAGE == "".join(temp_stdout.getvalue()[:-1].
    ...                  splitlines(True)[1:])
    True

    Supply too many arguments on the command line:

    >>> save_stdout = sys.stdout
    >>> temp_stdout = StringIO()
    >>> sys.stdout = temp_stdout
    >>> sys.argv = ["", "", "",]
    >>> _config_path = config_path()
    >>> sys.stdout = save_stdout
    >>> USAGE == "".join(temp_stdout.getvalue()[:-1].
    ...                  splitlines(True)[1:])
    True

    Supply non-file argument:

    >>> save_stdout = sys.stdout
    >>> temp_stdout = StringIO()
    >>> sys.stdout = temp_stdout
    >>> _config_path = os.path.join(
    ...                    os.path.dirname(
    ...                        os.path.abspath(__file__)),
    ...                    TEST_PATH)
    >>> sys.argv = ["", _config_path]
    >>> _config_path = config_path()
    >>> sys.stdout = save_stdout
    >>> USAGE == "".join(temp_stdout.getvalue()[:-1].
    ...                  splitlines(True)[1:])
    True

    Supply nonexistent file argument:

    >>> save_stdout = sys.stdout
    >>> temp_stdout = StringIO()
    >>> sys.stdout = temp_stdout
    >>> _config_path = os.path.join(
    ...                    os.path.dirname(
    ...                        os.path.abspath(__file__)),
    ...                    TEST_PATH, "xxxxx")
    >>> sys.argv = ["", _config_path]
    >>> _config_path = config_path()
    >>> sys.stdout = save_stdout
    >>> USAGE == "".join(temp_stdout.getvalue()[:-1].
    ...                  splitlines(True)[1:])
    True

    Supply valid config path argument:

    >>> _config_path = os.path.join(
    ...                    os.path.dirname(
    ...                        os.path.abspath(__file__)),
    ...                    TEST_PATH, "config.py")
    >>> sys.argv = ["", _config_path]
    >>> _config_path == config_path()
    True
    """

    path = None
    try:
        if len(sys.argv) == 2:
            if sys.argv[1] == "-t" or sys.argv[1] == "--test":
                _test()
            else:
                path = sys.argv[1]
                if not os.path.exists(path):
                    raise IOError(path + \
                                  " does not exist.")
                elif not os.path.isfile(path):
                    raise IOError(path + \
                                  " is not a file.")
        else:
            raise IOError("Incorrect number of arguments supplied.")
    except IOError as exc:
        print exc
        print USAGE
    return path


def config(path):
    """
    Return the configuration from a file.

    Execute empty configuration:

    >>> _config_path = os.path.join(
    ...                    os.path.dirname(
    ...                        os.path.abspath(__file__)),
    ...                    TEST_PATH, "empty_config.py")
    >>> ncdir,csvdir, \
            ncfile_pattern,csvfile_pattern, \
            location,platform,packages = \
            config(_config_path)
    >>> ncdir
    >>> csvdir
    >>> ncfile_pattern
    >>> csvfile_pattern
    >>> location
    >>> platform
    >>> packages

    Execute nonexistent configuration:

    >>> save_stdout = sys.stdout
    >>> temp_stdout = StringIO()
    >>> sys.stdout = temp_stdout
    >>> _config_path = os.path.join(
    ...                    os.path.dirname(
    ...                        os.path.abspath(__file__)),
    ...                    TEST_PATH, "xxxxx")
    >>> ncdir,csvdir, \
            ncfile_pattern,csvfile_pattern, \
            location,platform,packages = \
            config(_config_path)
    >>> sys.stdout = save_stdout
    >>> USAGE == "".join(temp_stdout.getvalue()[:-1].
    ...                  splitlines(True)[1:])
    True
    >>> ncdir
    >>> csvdir
    >>> ncfile_pattern
    >>> csvfile_pattern
    >>> location
    >>> platform
    >>> packages

    Execute bad configuration:

    >>> save_stdout = sys.stdout
    >>> temp_stdout = StringIO()
    >>> sys.stdout = temp_stdout
    >>> _config_path = os.path.join(
    ...                    os.path.dirname(
    ...                        os.path.abspath(__file__)),
    ...                    TEST_PATH, "bad_config.py")
    >>> ncdir,csvdir, \
            ncfile_pattern,csvfile_pattern, \
            location,platform,packages = \
            config(_config_path)
    >>> sys.stdout = save_stdout
    >>> USAGE == "".join(temp_stdout.getvalue()[:-1].
    ...                  splitlines(True)[1:])
    True
    >>> ncdir
    >>> csvdir
    >>> ncfile_pattern
    >>> csvfile_pattern
    >>> location
    >>> platform
    >>> packages

    Execute valid configuration:

    >>> _config_path = os.path.join(
    ...                    os.path.dirname(
    ...                        os.path.abspath(__file__)),
    ...                    TEST_PATH, "config.py")
    >>> ncdir,csvdir, \
            ncfile_pattern,csvfile_pattern, \
            location,platform,packages = \
            config(_config_path)
    >>> ncdir == os.path.join(os.path.dirname(os.path.abspath(__file__)),
    ...                       TEST_PATH, "nc")
    True
    >>> csvdir == os.path.join(os.path.dirname(os.path.abspath(__file__)),
    ...                        TEST_PATH, "csv")
    True
    >>> ncfile_pattern == "%(location)s_%(platform)s_%(package)s_" \
                          "[0-9][0-9][0-9][0-9]_[0-9][0-9].nc"
    True
    >>> csvfile_pattern == "%(location)s_%(platform)s_%(month)s.csv"
    True
    >>> location == "location"
    True
    >>> platform == "platform"
    True
    >>> packages == (('system', ['voltage',
    ...                          'memory',
    ...                          'interval', ]),
    ...              ('turbidity', ['turbidity', ]),
    ...              ('optode_127', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('optode_167', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('optode_169', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('optode_170', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('optode_171', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('optode_172', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('optode_173', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('optode_175', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('optode_176', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('optode_178', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('optode_179', ['o2concentration',
    ...                              'airsaturation',
    ...                              'temperature',
    ...                              'calphase',
    ...                              'tcphase',
    ...                              'c1rph',
    ...                              'c2rph',
    ...                              'c1amp',
    ...                              'c2amp',
    ...                              'rawtemp', ]),
    ...              ('conductivity', ['conductivity',
    ...                                'temperature', ]),
    ...              ('pressure', ['pressure',
    ...                            'temperature', ]),
    ...              ('current', ['abs_speed',
    ...                           'direction',
    ...                           'v',
    ...                           'u',
    ...                           'heading',
    ...                           'tiltx',
    ...                           'tilty',
    ...                           'std_speed',
    ...                           'strength',
    ...                           'pings', ]),
    ...             )
    True
    """

    namespace = {}
    try:
        execfile(path, globals(), namespace)
    except IOError as exc:
        print exc
        print USAGE
    except SyntaxError as exc:
        print exc
        print USAGE
    return (namespace.get("NCDIR"),
            namespace.get("CSVDIR"),
            namespace.get("NCFILE_PATTERN"),
            namespace.get("CSVFILE_PATTERN"),
            namespace.get("LOCATION"),
            namespace.get("PLATFORM"),
            namespace.get("PACKAGES"),
           )


def create(csvpath, ncdir, ncfile_pattern,
           location, platform, month, packages):
    """
    Create a specific month of CSV spreadsheet from sponge data NetCDF files.
    """

    columns = ["time", ]
    varmap = [("time", None), ]

    # Get all the column headings and variable arrays
    for package, variables in packages:
        ncpath = ncfile_pattern % {"location": location,
                                   "platform": platform,
                                   "package": package, }
        ncpath = ncpath[:-34] + month + ".nc"
        ncpath = os.path.join(ncdir, ncpath)
        try:
            variables.append("time")
            cdf = CDF(ncpath)
            for variable in variables:
                var = ".".join([package, variable])
                try:
                    varmap.append((var, cdf.var(variable).get()))
                    if variable != "time":
                        columns.append(var)
                except CDFError:
                    print "Bypassing variable", var
        except CDFError:
            print "Bypassing package", package
    varmap = dict(varmap)

    # Sanity check the time arrays
    time_list = [var for var in varmap.keys() if var.endswith(".time")]
    time_list = zip(time_list[:-1], time_list[1:])
    for time0, time1 in time_list:
        if any(varmap[time0] != varmap[time1]):
            raise CDFError("Time arrays are not equivalent.")
    varmap["time"] = varmap[time_list[0][0]]

    # Create the rows.
    rows = [[] for cell in varmap["time"]]
    for column in columns:
        for pos in range(len(rows)):
            cell = varmap[column][pos]
            if column == "time":
                cell = time.strftime("%m/%d/%Y %H:%M:%S",
                                     time.gmtime(cell))
            rows[pos].append(cell)

    # Output the CSV file.
    handle = open(csvpath, "wb")
    writer = csv.writer(handle)
    writer.writerow(columns)
    writer.writerows(rows)
    handle.close()

    return


def query(ncdir, csvdir,
          ncfile_pattern, csvfile_pattern,
          location, platform, packages):
    """
    Create CSV spreadsheets from sponge data NetCDF files.
    """

    # Find all monthly NetCDF files for location/platform
    ncfile_patterns = [(package,
                        ncfile_pattern % {"location":location,
                                          "platform":platform,
                                          "package":package, },
                       ) for package, variables in packages]
    paths = dict([(package,
                   [month
                    for month in glob.glob(os.path.join(ncdir, pattern))
                    if os.path.isfile(month)],
                  ) for package, pattern in ncfile_patterns])

    # Find all months for location/platform
    months = list(set([os.path.splitext(path)[0][-7:]
                       for package, variables in packages
                       for path in paths[package]]))
    months.sort()

    # Create a directory for monthly CSV files.
    if not os.path.exists(csvdir):
        os.mkdir(csvdir, 0755)
    elif not os.path.isdir(csvdir):
        raise IOError("CSV directory name " + \
                       csvdir + \
                       " exists and is not a directory.")

    # Create a CSV file for each month
    for month in months:
        path = csvfile_pattern % {"location": location,
                                  "platform": platform,
                                  "month": month}
        path = os.path.join(csvdir, path)
        if not os.path.exists(path):
            print "=" * 40
            print "Creating", path
            try:
                create(path, ncdir, ncfile_pattern,
                       location, platform, month, packages)
            except CDFError:
                print "Time array mismatch amongst", month, "packages."
        elif not os.path.isfile(path):
            raise IOError("CSV file name " + \
                          path + \
                          " exists and is not a file.")
        elif month == months[-1]:
            # Always create last month
            print "=" * 40
            print "Creating", path
            create(path + "new", ncdir, ncfile_pattern,
                   location, platform, month, packages)
            os.remove(path)
            os.rename(path + "new", path)

    return


def _main():
    """
    Run module as script.

    Supply incomplete config file:

    >>> save_stdout = sys.stdout
    >>> temp_stdout = StringIO()
    >>> sys.stdout = temp_stdout
    >>> _config_path = os.path.join(
    ...                    os.path.dirname(
    ...                        os.path.abspath(__file__)),
    ...                    TEST_PATH, "incomplete_config.py")
    >>> sys.argv = ["", _config_path]
    >>> _main()
    >>> sys.stdout = save_stdout
    >>> USAGE == "".join(temp_stdout.getvalue()[:-1].
    ...                  splitlines(True)[1:])
    True
    """

    _config_path = config_path()
    if _config_path:
        _config = config(_config_path)
        if any(_config):
            if all(_config):
                query(*_config)
            else:
                print "Config file", _config_path, \
                      "does not contain all necessary parameters."
                print USAGE

    return

if __name__ == "__main__":
    _main()