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 sys
import os
import glob
import csv
from pycdf import CDF, CDFError
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 == temp_stdout.getvalue()[:-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 == temp_stdout.getvalue()[:-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 == temp_stdout.getvalue()[:-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 == temp_stdout.getvalue()[:-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:
        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 == temp_stdout.getvalue()[:-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 == temp_stdout.getvalue()[:-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 = {}
    namespace["NCDIR"] = None
    namespace["CSVDIR"] = None
    namespace["NCFILE_PATTERN"] = None
    namespace["CSVFILE_PATTERN"] = None
    namespace["LOCATION"] = None
    namespace["PLATFORM"] = None
    namespace["PACKAGES"] = None
    try:
        execfile(path, globals(), namespace)
    except IOError:
        print USAGE
    except SyntaxError:
        print USAGE
    return (namespace["NCDIR"],
            namespace["CSVDIR"],
            namespace["NCFILE_PATTERN"],
            namespace["CSVFILE_PATTERN"],
            namespace["LOCATION"],
            namespace["PLATFORM"],
            namespace["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", ]
    columns.extend([".".join([package, variable])
                    for package, variables in packages
                    for variable in variables])

    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:
            cdf = CDF(ncpath)
            for variable in variables:
                cdf.var(variable).get()
        except CDFError:
            pass

    handle = open(csvpath, "wb")
    writer = csv.DictWriter(handle, columns)
    headers = dict([(column,column) for column in columns])
    writer.writerow(headers)
    handle.close()

    return columns


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):
            create(path, ncdir, ncfile_pattern,
                   location, platform, 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
            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 == temp_stdout.getvalue()[:-1]
    True
    """

    _config_path = config_path()
    if _config_path:
        _config = config(_config_path)
        if all(_config):
            query(*_config)
        else:
            print USAGE

    return

if __name__ == "__main__":
    _main()