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 output file that ends with `.csv` or `.json`.",
file=sys.stderr,
)
def write_to_json(results=limit(cas), filename='data/cad_write.json'):
"""Write an iterable of `CloseApproach` objects to a JSON file.
The precise output specification is in `README.md`. Roughly, the output is
a list containing dictionaries, each mapping `CloseApproach` attributes to
their values and the 'neo' key mapping to a dictionary of the associated
NEO's attributes.
:param results: An iterable of `CloseApproach` objects.
:param filename: A Path-like object pointing to where the data should be
saved.
"""
with open(filename, 'w') as outfile:
daters = []
for ca in results:
neo_dict = {'datetime_utc': ca.time_str,
'distance_au': ca.distance,
'velocity_km_s': ca.velocity,
"neo": {
'designation': ca.neo.designation,
'name': ca.neo.name,
'diameter_km': ca.neo.diameter,
'potentially_hazardous': ca.neo.hazardous
}
}
daters.append(neo_dict)
json.dump(daters, outfile)
def write_to_csv(results=limit(cas), filename='data/neos_write.csv'):
"""Write an iterable of `CloseApproach` objects to a CSV file.
The precise output specification is in `README.md`. Roughly, each output
row corresponds to the information in a single close approach from the
`results` stream and its associated near-Earth object.
:param results: An iterable of `CloseApproach` objects.
:param filename: A Path-like object pointing to where the data should be
saved.
"""
fieldnames = ('datetime_utc', 'distance_au', 'velocity_km_s',
'designation', 'name', 'diameter_km',
'potentially_hazardous')
with open(filename, 'w') as outfile:
writer = csv.writer(outfile)
writer.writerow(fieldnames)
for row in results:
writer.writerow((row.time,
row.distance,
row.velocity,
row.neo.designation,
row.neo.name,
row.neo.diameter,
row.neo.hazardous))
def query(database, args):
"""Perform the `query` subcommand.
: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)
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):
try:
print(result)
except StopIteration:
break
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)
def test_limit_produces_an_iterable(self):
self.assertIsInstance(limit(self.iterable, 3),
collections.abc.Iterable)
self.assertIsInstance(limit(self.iterable, 5),
collections.abc.Iterable)
self.assertIsInstance(limit(self.iterable, 10),
collections.abc.Iterable)
self.assertIsInstance(limit(self.iterable), collections.abc.Iterable)
self.assertIsInstance(limit(self.iterable, 0),
collections.abc.Iterable)
self.assertIsInstance(limit(self.iterable, None),
collections.abc.Iterable)
def test_limit_iterator_without_limit(self):
self.assertEqual(tuple(limit(iter(self.iterable))), (0, 1, 2, 3, 4))
self.assertEqual(tuple(limit(iter(self.iterable), 0)), (0, 1, 2, 3, 4))
self.assertEqual(tuple(limit(iter(self.iterable), None)),
(0, 1, 2, 3, 4))
def test_limit_iterator_with_larger_limit(self):
self.assertEqual(tuple(limit(iter(self.iterable), 10)),
(0, 1, 2, 3, 4))
def test_limit_iterator_with_matching_limit(self):
self.assertEqual(tuple(limit(iter(self.iterable), 5)), (0, 1, 2, 3, 4))
def test_limit_iterator_with_smaller_limit(self):
self.assertEqual(tuple(limit(iter(self.iterable), 3)), (0, 1, 2))
def test_limit_iterable_with_limit(self):
self.assertEqual(tuple(limit(self.iterable, 3)), (0, 1, 2))
