From e4882d70afcad41f9fcb4578078d448685ab540c Mon Sep 17 00:00:00 2001
From: Jylpah <jylpah@gmail.com>
Date: Tue, 11 May 2021 20:55:13 +0300
Subject: [PATCH] export replays(): --replays option to export replays

---
 analyze_wotb_replays.py | 33 +++++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/analyze_wotb_replays.py b/analyze_wotb_replays.py
index 789d999..6de4f14 100755
--- a/analyze_wotb_replays.py
+++ b/analyze_wotb_replays.py
@@ -28,6 +28,8 @@ REPLAY_I = 0
 STAT_TANK_BATTLE_MIN = 100
 BATTLE_TIME_BUCKET = 3600*24*14
 
+OPT_EXPORT_REPLAY_DIR = 'export_replays'
+
 RE_SRC_IS_DB = re.compile(r'^DB:')
 
 class StatFunc:
@@ -1419,6 +1421,7 @@ async def main(argv):
 		parser.add_argument('--mapfile', type=str, default='maps.json', help='JSON file to read Blitz map names from. Default is "maps.json"')
 		parser.add_argument('--json', action='store_true', default=False, help='Export data in JSON format')
 		parser.add_argument('--csv', action='store_true', default=False, help='Export data in CSV format')
+		parser.add_argument('--replays', action='store_true', default=False, help='Export replay info in JSON format')
 		parser.add_argument('-o','--outfile', type=str, default='-', metavar="OUTPUT", help='File to write results. Default STDOUT')
 		parser.add_argument('--db', action='store_true', default=OPT_DB, help='Use DB - You are unlikely to have it')
 		parser.add_argument('--filters', type=str, default=None, help='Filter replays based on categories. Filters given in JSON format.\nUse array "[]" for multiple filters/values. see --mode help.\nExample: : [ {"tier" : [8,9,10] }, { "player_wins" : 5 }]')
@@ -1568,6 +1571,9 @@ async def main(argv):
 					csvwriter.writerows(rows)
 					bu.verbose_std('Results exported to: ' + export_file)			
 			
+			if args.replays:
+				await export_replays(replays)
+
 			bu.debug('Finished. Cleaning up..................')
 		except Exception as err:
 			bu.error(exception=err)
@@ -1634,6 +1640,25 @@ async def help_extended(db : motor.motor_asyncio.AsyncIOMotorDatabase = None, pa
 	print('No, I am not planning to fix these for now.')	
 
 
+async def export_replays(replays: dict):
+	"""Export replays"""
+	try:
+		replay_dir = OPT_EXPORT_REPLAY_DIR +  '_' + bu.get_date_str()
+		bu.verbose_std('Exporting replays to ' + replay_dir)
+		if not os.path.isdir(replay_dir):
+			os.mkdir(replay_dir)
+
+		i = 0
+		for replay in replays.values():
+			filename = wg.get_replay_filename(replay)
+			async with aiofiles.open(os.path.join(replay_dir, filename), 'w', encoding="utf8") as f:
+				await f.write(json.dumps(replay, indent=4))
+			i += 1
+		bu.verbose_std( str(i) + ' replays exported')
+	except Exception as err:
+		bu.error(exception=err)
+
+
 ## move the class PlayerHistogram?
 def set_histogram_buckets(json: dict):
 	global histogram_fields
@@ -1785,14 +1810,14 @@ def filter_results(results: list, filter_json : str) -> bool:
 	return None
 
 
-def filter_replays(replays: dict, results: set) -> list:
+def filter_replays(replays_in: dict, results: set) -> list:
 	"""Filter source replays based on (filtered) results"""
 	try:
-		ret_replays = dict()
+		replays = dict()
 		for res in results:
 			_id = res['_id']
-			ret_replays[_id] = replays[_id]
-		return ret_replays
+			replays[_id] = replays_in[_id]
+		return replays
 	except Exception as err:
 		bu.error(exception=err)
 	return None