Added locations to the snapshot

2 files changed, 23 insertions, 2 deletions

diff --git a/salvager/esi.py b/salvager/esi.py
index 4783e12..d6acf30 100644
--- a/salvager/esi.py
+++ b/salvager/esi.py
@@ -32,3 +32,7 @@ def _trim_killmail(km):
 
 killmail = _make_endpoint("killmails/{}/{}/", trim=_trim_killmail)
 type = _make_endpoint("universe/types/{}/", include=['type_id', 'name', 'group_id'])
+system = _make_endpoint(
+	"universe/systems/{}/", include=['system_id', 'constellation_id', 'name', 'position', 'security_status'])
+constellation = _make_endpoint("universe/constellations/{}/", include=['region_id', 'name'])
+region = _make_endpoint("universe/regions/{}/", include=['name'])
diff --git a/scrap.py b/scrap.py
index 2f81aa0..0a5c7d9 100644
--- a/scrap.py
+++ b/scrap.py
@@ -32,11 +32,27 @@ def expand_details(snapshot):
 	return snapshot
 
 
-def get_ships(snapshot):
+def get_types(snapshot):
 	ships = {x['victim']['ship_type_id'] for x in snapshot['killmails']}
 	return [esi.type(x) for x in ships]
 
 
+def get_locations(snapshot):
+	solar_system_ids = {x['solar_system_id'] for x in snapshot['killmails']}
+	solar_systems = [esi.system(x) for x in solar_system_ids]
+	constellation_ids = {x['constellation_id'] for x in solar_systems}
+	constellations = {x: esi.constellation(x) for x in constellation_ids}
+	region_ids = {x['region_id'] for x in constellations.values()}
+	regions = {x: esi.region(x) for x in region_ids}
+	for system in solar_systems:
+		constellation = constellations[system['constellation_id']]
+		region_id = constellation['region_id']
+		system['constellation_name'] = constellation['name']
+		system['region_name'] = regions[region_id]['name']
+		system['region_id'] = region_id
+	return solar_systems
+
+
 def output_name(args):
 	"""
 	Generates name of the output file based on the CLI *args*.
@@ -58,7 +74,8 @@ def main():
 	snapshot = zkill.parse_battle_report(args.url)
 	expand_hashes(snapshot)
 	expand_details(snapshot)
-	snapshot['ships'] = get_ships(snapshot)
+	snapshot['types'] = get_types(snapshot)
+	snapshot['locations'] = get_locations(snapshot)
 	filename = output_name(args)
 	with open(filename, 'w') as fd:
 		opts = {'indent': 4} if args.pretty else {}