main.py

#!/usr/bin/env python3
"""Explore a dataset of near-Earth objects and their close approaches to Earth.

See `README.md` for a detailed discussion of this project.

This script can be invoked from the command line::

    $ python3 main.py {inspect,query,interactive} [args]

The `inspect` subcommand looks up an NEO by name or by primary designation, and
optionally lists all of that NEO's known close approaches:

    $ python3 main.py inspect --pdes 1P
    $ python3 main.py inspect --name Halley
    $ python3 main.py inspect --verbose --name Halley

The `query` subcommand searches for close approaches that match given criteria:

    $ python3 main.py query --date 1969-07-29
    $ python3 main.py query --start-date 2020-01-01 --end-date 2020-01-31 --max-distance 0.025
    $ python3 main.py query --start-date 2050-01-01 --min-distance 0.2 --min-velocity 50
    $ python3 main.py query --date 2020-03-14 --max-velocity 25 --min-diameter 0.5 --hazardous
    $ python3 main.py query --start-date 2000-01-01 --max-diameter 0.1 --not-hazardous
    $ python3 main.py query --hazardous --max-distance 0.05 --min-velocity 30

The set of results can be limited in size and/or saved to an output file in CSV
or JSON format:

    $ python3 main.py query --limit 5 --outfile results.csv
    $ python3 main.py query --limit 15 --outfile results.json

The `interactive` subcommand loads the NEO database and spawns an interactive
command shell that can repeatedly execute `inspect` and `query` commands without
having to wait to reload the database each time. However, it doesn't hot-reload.

If needed, the script can load data from data files other than the default with
`--neofile` or `--cadfile`.
"""
import argparse
import cmd
import datetime
import pathlib
import shlex
import sys
import time

from extract import load_neos, load_approaches
from database import NEODatabase
from filters import create_filters, limit
from write import write_to_csv, write_to_json


# Paths to the root of the project and the `data` subfolder.
PROJECT_ROOT = pathlib.Path(__file__).parent.resolve()
DATA_ROOT = PROJECT_ROOT / 'data'

# The current time, for use with the kill-on-change feature of the interactive shell.
_START = time.time()


def date_fromisoformat(date_string):
    """Return a `datetime.date` corresponding to a string in YYYY-MM-DD format.

    In Python 3.7+, there is `datetime.date.fromisoformat`, but alas - we're
    supporting Python 3.6+.

    :param date_string: A date in the format YYYY-MM-DD.
    :return: A `datetime.date` correspondingo the given date string.
    """
    try:
        return datetime.datetime.strptime(date_string, '%Y-%m-%d').date()
    except ValueError:
        raise argparse.ArgumentTypeError(f"'{date_string}' is not a valid date. Use YYYY-MM-DD.")


def make_parser():
    """Create an ArgumentParser for this script.

    :return: A tuple of the top-level, inspect, and query parsers.
    """
    parser = argparse.ArgumentParser(
        description="Explore past and future close approaches of near-Earth objects."
    )

    # Add arguments for custom data files.
    parser.add_argument('--neofile', default=(DATA_ROOT / 'neos.csv'),
                        type=pathlib.Path,
                        help="Path to CSV file of near-Earth objects.")
    parser.add_argument('--cadfile', default=(DATA_ROOT / 'cad.json'),
                        type=pathlib.Path,
                        help="Path to JSON file of close approach data.")
    subparsers = parser.add_subparsers(dest='cmd')

    # Add the `inspect` subcommand parser.
    inspect = subparsers.add_parser('inspect',
                                    description="Inspect an NEO by primary designation or by name.")
    inspect.add_argument('-v', '--verbose', action='store_true',
                         help="Additionally, print all known close approaches of this NEO.")
    inspect_id = inspect.add_mutually_exclusive_group(required=True)
    inspect_id.add_argument('-p', '--pdes',
                            help="The primary designation of the NEO to inspect (e.g. '433').")
    inspect_id.add_argument('-n', '--name',
                            help="The IAU name of the NEO to inspect (e.g. 'Halley').")

    # Add the `query` subcommand parser.
    query = subparsers.add_parser('query',
                                  description="Query for close approaches that "
                                              "match a collection of filters.")
    filters = query.add_argument_group('Filters',
                                       description="Filter close approaches by their attributes "
                                                   "or the attributes of their NEOs.")
    filters.add_argument('-d', '--date', type=date_fromisoformat,
                         help="Only return close approaches on the given date, "
                              "in YYYY-MM-DD format (e.g. 2020-12-31).")
    filters.add_argument('-s', '--start-date', type=date_fromisoformat,
                         help="Only return close approaches on or after the given date, "
                              "in YYYY-MM-DD format (e.g. 2020-12-31).")
    filters.add_argument('-e', '--end-date', type=date_fromisoformat,
                         help="Only return close approaches on or before the given date, "
                              "in YYYY-MM-DD format (e.g. 2020-12-31).")
    filters.add_argument('--min-distance', dest='distance_min', type=float,
                         help="In astronomical units. Only return close approaches that "
                              "pass as far or farther away from Earth as the given distance.")
    filters.add_argument('--max-distance', dest='distance_max', type=float,
                         help="In astronomical units. Only return close approaches that "
                              "pass as near or nearer to Earth as the given distance.")
    filters.add_argument('--min-velocity', dest='velocity_min', type=float,
                         help="In kilometers per second. Only return close approaches "
                              "whose relative velocity to Earth at approach is as fast or faster "
                              "than the given velocity.")
    filters.add_argument('--max-velocity', dest='velocity_max', type=float,
                         help="In kilometers per second. Only return close approaches "
                              "whose relative velocity to Earth at approach is as slow or slower "
                              "than the given velocity.")
    filters.add_argument('--min-diameter', dest='diameter_min', type=float,
                         help="In kilometers. Only return close approaches of NEOs with "
                              "diameters as large or larger than the given size.")
    filters.add_argument('--max-diameter', dest='diameter_max', type=float,
                         help="In kilometers. Only return close approaches of NEOs with "
                              "diameters as small or smaller than the given size.")
    filters.add_argument('--hazardous', dest='hazardous', default=None, action='store_true',
                         help="If specified, only return close approaches of NEOs that "
                              "are potentially hazardous.")
    filters.add_argument('--not-hazardous', dest='hazardous', default=None, action='store_false',
                         help="If specified, only return close approaches of NEOs that "
                              "are not potentially hazardous.")
    query.add_argument('-l', '--limit', type=int,
                       help="The maximum number of matches to return. "
                            "Defaults to 10 if no --outfile is given.")
    query.add_argument('-o', '--outfile', type=pathlib.Path,
                       help="File in which to save structured results. "
                            "If omitted, results are printed to standard output.")

    repl = subparsers.add_parser('interactive',
                                 description="Start an interactive command session "
                                             "to repeatedly run `interact` and `query` commands.")
    repl.add_argument('-a', '--aggressive', action='store_true',
                      help="If specified, kill the session whenever a project file is modified.")
    return parser, inspect, query


def inspect(database, pdes=None, name=None, verbose=False):
    """Perform the `inspect` subcommand.

    This function fetches an NEO by designation or by name. If a matching NEO is
    found, information about the NEO is printed (additionally, information for
    all of the NEO's known close approaches is printed if `verbose=True`).
    Otherwise, a message is printed noting that there are no matching NEOs.

    At least one of `pdes` and `name` must be given. If both are given, prefer
    to look up the NEO by the primary designation.

    :param database: The `NEODatabase` containing data on NEOs and their close approaches.
    :param pdes: The primary designation of an NEO for which to search.
    :param name: The name of an NEO for which to search.
    :param verbose: Whether to additionally print all of a matching NEO's close approaches.
    :return: The matching `NearEarthObject`, or None if not found.
    """
    # Fetch the NEO of interest.
    if pdes:
        neo = database.get_neo_by_designation(pdes)
    else:
        neo = database.get_neo_by_name(name)

    # Ensure that we have received an NEO.
    if not neo:
        print("No matching NEOs exist in the database.", file=sys.stderr)
        return None

    # Display information about this NEO, and optionally its close approaches if verbose.
    print(neo)
    if verbose:
        for approach in neo.approaches:
            print(f"- {approach}")
    return neo


def query(database, args):
    """Perform the `query` subcommand.

    Create a collection of filters with `create_filters` and supply them to the
    database's `query` method to produce a stream of matching results.

    If an output file wasn't given, print these results to stdout, limiting to
    10 entries if no limit was specified. If an output file was given, use the
    file's extension to infer whether the file should hold CSV or JSON data, and
    then write the results to the output file in that format.

    :param database: The `NEODatabase` containing data on NEOs and their close approaches.
    :param args: All arguments from the command line, as parsed by the top-level parser.
    """
    # Construct a collection of filters from arguments supplied at the command line.
    filters = create_filters(
        date=args.date, start_date=args.start_date, end_date=args.end_date,
        distance_min=args.distance_min, distance_max=args.distance_max,
        velocity_min=args.velocity_min, velocity_max=args.velocity_max,
        diameter_min=args.diameter_min, diameter_max=args.diameter_max,
        hazardous=args.hazardous
    )
    # Query the database with the collection of filters.
    results = database.query(filters)

    if not args.outfile:
        # Write the results to stdout, limiting to 10 entries if not specified.
        for result in limit(results, args.limit or 10):
            print(result)
    else:
        # Write the results to a file.
        if args.outfile.suffix == '.csv':
            write_to_csv(limit(results, args.limit), args.outfile)
        elif args.outfile.suffix == '.json':
            write_to_json(limit(results, args.limit), args.outfile)
        else:
            print("Please use an output file that ends with `.csv` or `.json`.", file=sys.stderr)


class NEOShell(cmd.Cmd):
    """Perform the `interactive` subcommand.

    This is a `cmd.Cmd` shell - a specialized tool for command-based REPL sessions.

    It wraps the `inspect` and `query` parsers to parse flags for those commands
    as if they were supplied at the command line.

    The primary purpose of this shell is to allow users to repeatedly perform
    inspect and query commands, while only loading the data (which can be quite
    slow) once.
    """
    intro = ("Explore close approaches of near-Earth objects. "
             "Type `help` or `?` to list commands and `exit` to exit.\n")
    prompt = '(neo) '

    def __init__(self, database, inspect_parser, query_parser, aggressive=False, **kwargs):
        """Create a new `NEOShell`.

        Creating this object doesn't start the session - for that, use `.cmdloop()`.

        :param database: The `NEODatabase` containing data on NEOs and their close approaches.
        :param inspect_parser: The subparser for the `inspect` subcommand.
        :param query_parser: The subparser for the `query` subcommand.
        :param aggressive: Whether to kill the session whenever a project file is changed.
        :param kwargs: A dictionary of excess keyword arguments passed to the superclass.
        """
        super().__init__(**kwargs)
        self.db = database
        self.inspect = inspect_parser
        self.query = query_parser
        self.aggressive = aggressive

    @classmethod
    def parse_arg_with(cls, arg, parser):
        """Parse the additional text passed to a command, using a given parser.

        If any error is encountered (in lexical parsing or argument parsing),
        print the error to stderr and return None.

        :param arg: The additional text supplied after the command.
        :param parser: An `argparse.ArgumentParser` to parse the arguments.
        :return: A `Namespace` of the arguments (produced by `parse_args`) or None.
        """
        # Lexically parse the additional text with POSIX shell-like syntax.
        try:
            args = shlex.split(arg)
        except ValueError as err:
            print(err, file=sys.stderr)
            return None

        # Use the ArgumentParser to parse the shell arguments.
        try:
            return parser.parse_args(args)
        except SystemExit as err:
            # The `parse_args` method doesn't actually surface `ArgumentError`s
            # nor `ArgumentTypeError`s - instead, it calls its own `error`
            # method which prints the error message and then calls `sys.exit`.
            return None

    def do_i(self, arg):
        """Shorthand for `inspect`."""
        self.do_inspect(arg)

    def do_inspect(self, arg):
        """Perform the `inspect` subcommand within the REPL session.

        Inspect an NEO by designation or by name:

            (neo) inspect --pdes 1P
            (neo) inspect --name Halley

        Additionally, list all known close approaches:

            (neo) inspect --verbose --name Eros
        """
        args = self.parse_arg_with(arg, self.inspect)
        if not args:
            return

        # Run the `inspect` subcommand.
        inspect(self.db,
                pdes=args.pdes, name=args.name,
                verbose=args.verbose)

    def do_q(self, arg):
        """Shorthand for `query`."""
        self.do_query(arg)

    def do_query(self, arg):
        """Perform the `query` subcommand within the REPL session.

        This command behaves the same as the `query` subcommand from the command
        line. For example, to query close approaches on January 1st, 2020:

            (neo) query --date 2020-01-01

        You can use any of the other filters: `--start-date`, `--end-date`,
        `--min-distance`, `--max-distance`, `--min-velocity`, `--max-velocity`,
        `--min-diameter`, `--max-diameter`, `--hazardous`, `--not-hazardous`.

        The number of results shown can be limited to a maximum number with `--limit`:

            (neo) query --limit 2

        The results can be saved to a file (instead of displayed to stdout) with
        `--outfile`:

            (neo) query --limit 5 --outfile results.csv
            (neo) query --limit 5 --outfile results.json
        """
        args = self.parse_arg_with(arg, self.query)
        if not args:
            return

        # Run the `inspect` subcommand.
        query(self.db, args)

    def do_EOF(self, _arg):
        """Exit the interactive session."""
        return True

    # Alternative ways to quit.
    do_exit = do_EOF
    do_quit = do_EOF

    def precmd(self, line):
        """Watch for changes to the files in this project."""
        changed = [f for f in PROJECT_ROOT.glob('*.py') if f.stat().st_mtime > _START]
        if changed:
            print("The following file(s) have been modified since this interactive session began: "
                  f"{', '.join(str(f.relative_to(PROJECT_ROOT)) for f in changed)}.",
                  file=sys.stderr)
            if not self.aggressive:
                print("To include these changes, please exit and restart this interactive session.",
                      file=sys.stderr)
            else:
                print("Preemptively terminating the session aggressively.", file=sys.stderr)
                return 'exit'
        return line


def main():
    """Run the main script."""
    parser, inspect_parser, query_parser = make_parser()
    args = parser.parse_args()

    # Extract data from the data files into structured Python objects.
    database = NEODatabase(load_neos(args.neofile), load_approaches(args.cadfile))

    # Run the chosen subcommand.
    if args.cmd == 'inspect':
        inspect(database, pdes=args.pdes, name=args.name, verbose=args.verbose)
    elif args.cmd == 'query':
        query(database, args)
    elif args.cmd == 'interactive':
        NEOShell(database, inspect_parser, query_parser, aggressive=args.aggressive).cmdloop()


if __name__ == '__main__':
    main()