Merging from branch nil, which shouldn't exist

src/plotypus/cli.py

@@ -17,13 +17,22 @@
 from plotypus.periodogram import Lomb_Scargle, conditional_entropy
 import plotypus
 from plotypus.preprocessing import Fourier
-from plotypus.utils import mad, make_sure_path_exists, pmap, verbose_print
+from plotypus.utils import (mad, make_sure_path_exists, pmap, valid_basename,
+                            verbose_print)
 from plotypus.resources import matplotlibrc
 
 import pkg_resources # part of setuptools
 __version__ = pkg_resources.require("plotypus")[0].version
 
 
+# Dict mapping higher level output types (table, plot) to the specific
+# quantaties output (lightcurve, residual, etc). Any new output types need to
+# be added here.
+outputs = {
+    "table": ["lightcurve", "residual"],
+    "plot":  ["lightcurve", "residual"]
+}
+
 def get_args():
     parser = ArgumentParser()
     general_group    = parser.add_argument_group('General')
@@ -44,6 +53,26 @@ def get_args():
         default=None,
         help='location of stellar observations '
              '(default = stdin)')
+    general_group.add_argument('--output-all', type=str,
+        default=None, metavar="DIR",
+        help='(optional) location to output all file types (plots and tables '
+             'alike), creating a sub-directory for each star. individual '
+             'filenames have defaults which can be overridden by the options '
+             'of the same name. '
+             'cannot be used in conjunction with --output-{table,plot}-all '
+             'option')
+    general_group.add_argument('--output-table-all', type=str,
+        default=None, metavar="DIR",
+        help='(optional) location to output all table file types, creating a '
+             'sub-directory for each star. individual filenames have defaults '
+             'which can be overridden by the options of the same name. '
+             'cannot be used in conjunction with --output-all option')
+    general_group.add_argument('--output-plot-all', type=str,
+        default=None, metavar="DIR",
+        help='(optional) location to output all plot file types, creating a '
+             'sub-directory for each star. individual filenames have defaults '
+             'which can be overridden by the options of the same name. '
+             'cannot be used in conjunction with --output-all option')
     general_group.add_argument('-n', '--star-name', type=str,
         default=None,
         help='name of star '
@@ -116,16 +145,20 @@ def get_args():
 
     lightcurve_group.add_argument('--output-table-lightcurve', type=str,
         default=None,
-        help='(optional) location to save fitted light curve tables')
+        help='(optional) location to save fitted light curve tables, '
+             'or prefix to filename if any of `--output-*-all` options used.')
     lightcurve_group.add_argument('--output-plot-lightcurve', type=str,
         default=None,
-        help='(optional) location to save light curve plots')
+        help='(optional) location to save light curve plots, '
+             'or prefix to filename if any of `--output-*-all` options used.')
     lightcurve_group.add_argument('--output-table-residual', type=str,
         default=None,
-        help='(optional) location to save the fit residuals tables')
+        help='(optional) location to save the fit residuals tables, '
+             'or prefix to filename if any of `--output-*-all` options used.')
     lightcurve_group.add_argument('--output-plot-residual', type=str,
         default=None,
-        help='(optional) location to save the fit residuals plots')
+        help='(optional) location to save the fit residuals plots, '
+             'or prefix to filename if any of `--output-*-all` options used.')
     lightcurve_group.add_argument('--phase-points', type=int,
         default=100, metavar='N',
         help='number of phase points to output '
@@ -301,6 +334,8 @@ def get_args():
         args.parameters = read_table(args.parameters, args.param_sep,
                                      index_col=0, engine='python')
 
+    output_setup(args)
+
     return args
 
 
@@ -309,7 +344,7 @@ def main():
     args_dict = vars(args)
 
     min_degree, max_degree = args.fourier_degree
-    filenames = list(map(lambda x: x.strip(), _get_files(args.input)))
+    filenames = list(map(lambda x: x.strip(), get_files(args.input)))
     filepaths = map(lambda filename:
                     filename if path.isfile(filename)
                              else path.join(args.input, filename),
@@ -345,12 +380,138 @@ def main():
 
     def printer(result):
         if result is not None:
-            _print_star(result, max_degree, args.series_form,
-                        args.format, sep)
+            print_star(result, max_degree, args.series_form,
+                       args.format, sep)
     pmap(process_star, filepaths, callback=printer,
          processes=args.star_processes, **picklable_args)
 
 
+def fname_fn_bulk(directory, prefix):
+    """fname_fn_bulk(directory, prefix)
+    Returns a function which takes a name, and returns a filename in the given
+    `directory`, whose basename has the given `prefix` and `suffix`.
+    """
+    def fn(name, ext):
+        full_directory = path.join(directory, name)
+        make_sure_path_exists(full_directory)
+        return path.join(full_directory, prefix+ext)
+
+    return fn
+
+
+def fname_fn_individual(directory):
+    """fname_fn_individual(directory)
+    Returns a function which takes a name, and returns a filename in the given
+    `directory`, whose basename has the given `prefix` and `suffix`.
+    """
+    def fn(name, ext):
+        return path.join(directory, name+ext)
+
+    return fn
+
+
+def output_attr_name_from_parent_child(parent, child):
+    """output_attr_name_from_parent_child(parent, child)
+
+    Returns the attribute name of an output type. These attributes have names
+    of the form output_<parent>_<child>.
+
+    Examples include "output_plot_lightcurve" and "output_table_residuals".
+    """
+    return "_".join(["output", parent, child])
+
+
+def output_setup_bulk(args, directory, parent, child):
+    """output_setup_bulk(args, directory, parent, child)
+
+    Sets the output filename function for the given parent-child output type,
+    assuming this is a bulk output.
+    """
+    # string representing the name of the output type attribute
+    attr_name = output_attr_name_from_parent_child(parent, child)
+
+    # the attribute should either be the string representation of the filename
+    # prefix, or if one wasn't given, default to the name of the output type
+    prefix = getattr(args, attr_name, None)
+    if prefix is None:
+        prefix = child
+
+    # exit in error if prefix contains a directory or is the empty string
+    if not valid_basename(prefix) or prefix == "":
+        print(prefix + " is not a valid filename prefix",
+              file=stderr)
+        exit(1)
+
+    # replace the output type attribute with a function taking as input a star's
+    # name and file extension, and returning the desired filename string
+    setattr(args, attr_name, fname_fn_bulk(directory, prefix))
+
+
+def output_setup_individual(args, parent, child):
+    """output_setup_individual(args, parent, child)
+
+    Sets the output filename function for the given parent-child output type,
+    assuming this is an individual output.
+    """
+    # string representing the name of the output type attribute
+    attr_name = output_attr_name_from_parent_child(parent, child)
+    # the attribute should either be the string representation of the output
+    # directory, or None if one wasn't given
+    directory = getattr(args, attr_name, None)
+
+    # replace the output type attribute with a function taking as input a star's
+    # name and file extension, and returning the desired filename string,
+    # assuming a directory was given
+    if directory is not None:
+        make_sure_path_exists(directory)
+        setattr(args, attr_name, fname_fn_individual(directory))
+
+
+def output_setup(args):
+    """output_setup(args)
+
+    Sets the output filename functions for all of the various output types
+    requested.
+    """
+    # check for illegal combinations
+    if (args.output_all is not None and
+        any(x is not None for x in [args.output_table_all,
+                                    args.output_plot_all])):
+        print("Cannot combine --output-all with --output-table-all or",
+              "--output-plot-all flags.",
+              file=stderr)
+        exit(1)
+
+    # output all filetypes
+    if args.output_all is not None:
+        make_sure_path_exists(args.output_all)
+        for parent, children in outputs.items():
+            for child in children:
+                output_setup_bulk(args, args.output_all,
+                                  parent, child)
+
+    # output all table/plot filetypes to separate directories
+    if args.output_table_all is not None:
+        make_sure_path_exists(args.output_table_all)
+        for child in outputs["table"]:
+            output_setup_bulk(args, args.output_table_all,
+                              "table", child)
+    if args.output_plot_all is not None:
+        make_sure_path_exists(args.output_plot_all)
+        for child in outputs["plot"]:
+            output_setup_bulk(args, args.output_plot_all,
+                              "plot", child)
+
+    # output all individual filetypes explicitly requested, assuming none of
+    # the bulk operations above were selected
+    if all(x is None for x in [args.output_all,
+                               args.output_table_all,
+                               args.output_plot_all]):
+        for parent, children in outputs.items():
+            for child in children:
+                output_setup_individual(args, parent, child)
+
+
 def process_star(filename,
                  *,
                  extension, star_name, period, shift,
@@ -359,8 +520,10 @@ def process_star(filename,
                  output_table_lightcurve, output_plot_lightcurve,
                  output_table_residual,   output_plot_residual,
                  **kwargs):
-    """Processes a star's lightcurve, prints its coefficients, and saves
-    its plotted lightcurve to a file. Returns the result of get_lightcurve.
+    """process_star(filename, *, extension, star_name, period, shift, parameters, period_label, shift_label, plot_engine, output_table_lightcurve, output_plot_lightcurve, output_table_residual, output_plot_residual, **kwargs)
+
+    Processes a star's lightcurve, prints its coefficients, and saves its
+    plotted lightcurve to a file. Returns the result of get_lightcurve.
     """
     if star_name is None:
         basename = path.basename(filename)
@@ -390,11 +553,8 @@ def process_star(filename,
 
     # output the phased lightcurve as a table
     if output_table_lightcurve is not None:
-        # create the output path if non-existant
-        make_sure_path_exists(output_table_lightcurve)
         # construct the filename for the output table
-        filename = path.join(output_table_lightcurve,
-                             result["name"] + extension)
+        filename = output_table_lightcurve(result["name"], extension)
         # save the table to a file
         savetxt_lightcurve(filename, result["lightcurve"],
                            fmt=kwargs["format"],
@@ -406,7 +566,14 @@ def process_star(filename,
         # we wish to unpack.
         combined_kwargs = dict(kwargs)
         combined_kwargs.update(result)
-        plot = plot_lightcurve(output=output_plot_lightcurve,
+
+        # currenlty only .tikz output needs an explicit extension,
+        # mpl provides its own
+        extension = ".tikz" if plot_engine == "tikz" else ""
+        # construct the filename for the output plot
+        filename = output_plot_lightcurve(result["name"], extension)
+        # make the plot
+        plot = plot_lightcurve(output=filename,
                                engine=plot_engine,
                                **combined_kwargs)
         # allow figure to get garbage collected
@@ -419,20 +586,23 @@ def process_star(filename,
 
     # output the residuals as a table
     if output_table_residual is not None:
-        # create the output path if non-existant
-        make_sure_path_exists(output_table_residual)
         # construct the filename for the output table
-        filename = path.join(output_table_residual,
-                             result["name"] + extension)
+        filename = output_table_residual(result["name"], extension)
         # save the table to a file
         numpy.savetxt(filename, result["residuals"],
                       fmt=kwargs["format"],
                       delimiter=kwargs["output_sep"])
 
     # output the residuals as a plot
     if output_plot_residual is not None:
+        # currenlty only .tikz output needs an explicit extension,
+        # mpl provides its own
+        extension = ".tikz" if plot_engine == "tikz" else ""
+        # construct the filename for the output plot
+        filename = output_plot_residual(result["name"], extension)
+        # make the plot
         plot = plot_residual(result["name"], result["residuals"],
-                             output=output_plot_residual,
+                             output=filename,
                              sanitize_latex=kwargs["sanitize_latex"])
         # allow figure to get garbage collected
         if plot_engine == "mpl":
@@ -442,24 +612,27 @@ def process_star(filename,
             from matplotlib.pyplot import close
             close(plot)
 
-
     return result
 
 
-def _print_star(result, max_degree, form, fmt, sep):
+def print_star(result, max_degree, form, fmt, sep):
+    """print_star(result, max_degree, form, fmt, sep)
+
+    Prints a star's entry in the table printed to stdout.
+    """
     if result is None:
         return
 
     def format_one(x):
-        """format one number string according to fmt"""
+        """Format one number string according to fmt."""
         return fmt % x
 
     def format_all(x):
-        """format every number in a sequence according to fmt"""
+        """Format every number in a sequence according to fmt."""
         return map(format_one, x)
 
     def format_keys(*keys):
-        """format result[key] for the list of keys provided"""
+        """Format result[key] for the list of keys provided."""
         return format_all(result[key] for key in keys)
 
     # count inliers and outliers
@@ -500,11 +673,18 @@ def format_keys(*keys):
           sep=sep)
 
 
-def _get_files(input):
+def get_files(input):
+    """get_files(input)
+
+    Returns an iterable of filenames containing data to process.
+    """
     if input is None:
         return stdin
     elif input[0] == "@":
         with open(input[1:], 'r') as f:
+            # TODO:
+            # check if mapping over `f` instead of `f.readlines()` has same
+            # result
             return map(lambda x: x.strip(), f.readlines())
     elif path.isfile(input):
         return [input]