Added type details and moved esi stuff to own module

2 files changed, 43 insertions, 14 deletions

diff --git a/salvager/esi.py b/salvager/esi.py
new file mode 100644
index 0000000..4783e12
--- /dev/null
+++ b/salvager/esi.py
@@ -0,0 +1,34 @@
+import requests
+
+
+def _make_endpoint(path, include=None, exclude=None, trim=None):
+	def _trim(data):
+		if include:
+			data = {k: v for k, v in data.items() if k in include}
+		if exclude:
+			for key in exclude:
+				if key in data:
+					del data[key]
+		return data
+
+	if not callable(trim):
+		trim = _trim
+
+	def _get(*args):
+		query = "https://esi.evetech.net/latest/" + path + "?datasource=tranquility"
+		response = requests.get(query.format(*args))
+		response.raise_for_status()
+		return trim(response.json())
+
+	return _get
+
+
+def _trim_killmail(km):
+	del km['attackers']
+	del km['victim']['items']
+	del km['victim']['damage_taken']
+	return km
+
+
+killmail = _make_endpoint("killmails/{}/{}/", trim=_trim_killmail)
+type = _make_endpoint("universe/types/{}/", include=['type_id', 'name', 'group_id'])
diff --git a/scrap.py b/scrap.py
index dc10eb5..2939451 100644
--- a/scrap.py
+++ b/scrap.py
@@ -7,6 +7,8 @@ from html.parser import HTMLParser
 
 import requests
 
+from salvager import esi
+
 
 OWNER_HREF = re.compile(r"/(?:corporation|alliance)/(\d+)/?")
 SYSTEM_AND_DATE = re.compile(r"/(\d+)/(\d+)/?$")
@@ -90,30 +92,22 @@ def expand_hashes(snapshot):
 	return snapshot
 
 
-def get_details(kill, hash):
-	"""
-	Retrieves detailed information about killmail from EVE ESI using killmail's *kill* ID and *hash*.
-	"""
-	query = "https://esi.evetech.net/latest/killmails/{}/{}/?datasource=tranquility"
-	response = requests.get(query.format(kill, hash))
-	response.raise_for_status()
-	return response.json()
-
-
 def expand_details(snapshot):
 	"""
 	Expands killmails in *snapshot* IN PLACE by adding details from EVE ESI. Some data is dropped in process as e.g.,
 	full information on attackers is not important in context of the visualizations.
 	"""
 	for killmail in snapshot['killmails']:
-		details = get_details(killmail['killmail_id'], killmail['hash'])
-		del details['attackers']
-		del details['victim']['items']
-		del details['victim']['damage_taken']
+		details = esi.killmail(killmail['killmail_id'], killmail['hash'])
 		killmail.update(details)
 	return snapshot
 
 
+def get_ships(snapshot):
+	ships = {x['victim']['ship_type_id'] for x in snapshot['killmails']}
+	return [esi.type(x) for x in ships]
+
+
 def output_name(args):
 	"""
 	Generates name of the output file based on the CLI *args*.
@@ -135,6 +129,7 @@ def main():
 	snapshot = get_related_kills(args.url)
 	expand_hashes(snapshot)
 	expand_details(snapshot)
+	snapshot['ships'] = get_ships(snapshot)
 	filename = output_name(args)
 	with open(filename, 'w') as fd:
 		opts = {'indent': 4} if args.pretty else {}