extract.py

"""Extract data on near-Earth objects and close approaches from CSV and JSON files.

The `load_neos` function extracts NEO data from a CSV file, formatted as
described in the project instructions, into a collection of `NearEarthObject`s.

The `load_approaches` function extracts close approach data from a JSON file,
formatted as described in the project instructions, into a collection of
`CloseApproach` objects.

The main module calls these functions with the arguments provided at the command
line, and uses the resulting collections to build an `NEODatabase`.

You'll edit this file in Task 2.
"""
import csv
import json

from models import CloseApproach, NearEarthObject


def load_neos(neo_csv_path):
    """Read near-Earth object information from a CSV file.

    :param neo_csv_path: A path to a CSV file containing data about near-Earth objects.
    :return: A collection of `NearEarthObject`s.
    """
    neos_list = []
    with neo_csv_path.open("r") as infile:
        reader = csv.DictReader(infile)
        for row in reader:
            designation = row["pdes"]
            name = row["name"]
            hazardous = row["pha"]
            diameter = row["diameter"]
            neos_list.append(NearEarthObject(designation, name, hazardous, diameter))
    return neos_list


def load_approaches(cad_json_path):
    """Read close approach data from a JSON file.

    :param cad_json_path: A path to a JSON file containing data about close approaches.
    :return: A collection of `CloseApproach`es.
    """
    close_approach_list = []
    with cad_json_path.open("r") as infile:
        contents = json.load(infile)
        for approaches in contents["data"]:
            designation = approaches[contents["fields"].index("des")]
            time = approaches[contents["fields"].index("cd")]
            distance = approaches[contents["fields"].index("dist")]
            velocity = approaches[contents["fields"].index("v_rel")]
            close_approach_list.append(
                CloseApproach(time, designation, distance, velocity)
            )
    return close_approach_list