Add find_equip.py

README.md

@@ -342,15 +342,7 @@ To enable support for multiple users perform the steps below:
   * If multiplayer is enabled, access ``https://<zoffline_ip>/login/``, check "Enable ghosts" and click "Start Zwift" to save the option.
 * When you save an activity, the ghost will be saved in ``storage/<player_id>/ghosts/<world>/<route>``. Next time you ride the same route, the ghost will be loaded.
 * Type ``.regroup`` in chat to regroup the ghosts.
-* Equipment ("bike_frame", "bike_wheel_front", "bike_wheel_rear", "ride_jersey", "ride_helmet_type", "glasses_type", "ride_shoes_type" and "ride_socks_type") can be customized by creating a file ``ghost_profile.txt`` inside the ``storage`` folder. Values (signature) can be found in [GameDictionary](http://cdn.zwift.com/gameassets/GameDictionary.xml).
-
-  Example:
-  ``` json
-  {
-    "ride_jersey": 90433914,
-    "bike_frame": 1456463855
-  }
-  ```
+* Equipment can be customized by creating a file ``ghost_profile.txt`` inside the ``storage`` folder. The script ``find_equip.py`` can be used to populate this file.
 
 #### Bots
 

scripts/find_equip.py

@@ -0,0 +1,85 @@
+#! /usr/bin/env python
+
+# pip install country-converter fuzzywuzzy
+#
+# use to find and print entries from http://cdn.zwift.com/gameassets/GameDictionary.xml
+# scripts/find_equip.py -p emonda -s sworks -o blue -j zwift -e kask
+
+import urllib.request
+import json
+import country_converter as coco
+import argparse
+import os
+import sys
+import xml.etree.ElementTree as ET
+from fuzzywuzzy import process
+from fuzzywuzzy import fuzz
+
+cc = coco.CountryConverter()
+
+gd_file = 'GameDictionary.xml'
+if not os.path.isfile(gd_file):
+    open(gd_file, 'wb').write(urllib.request.urlopen('http://cdn.zwift.com/gameassets/%s' % gd_file).read())
+tree = ET.parse(gd_file)
+
+def get_item(location, item, json_name):
+    items = {}
+    for x in tree.findall(location):
+        items[x.get('name')] = int(x.get('signature'))
+    best_match = process.extractOne(item, items.keys(), scorer=fuzz.token_set_ratio)
+    equip = {}
+    equip[json_name] = items[best_match[0]]
+    equip[json_name+'_name'] = best_match[0]
+    return equip
+
+def main(argv):
+    global args
+
+    parser = argparse.ArgumentParser(description='Print out json for bot.txt or ghost_profile.txt')
+    megroup = parser.add_mutually_exclusive_group()
+    parser.add_argument('-n', '--nation', help='Specify nation', default=False)
+    parser.add_argument('-j', '--jersey', help='Get jersey', default=False)
+    megroup.add_argument('-b', '--bike', help='Get bike', default=False)
+    parser.add_argument('-w', '--wheels', help='Get wheels', default=False)
+    parser.add_argument('-e', '--helmet', help='Get helmet', default=False)
+    parser.add_argument('-s', '--shoes', help='Get shoes', default=False)
+    parser.add_argument('-o', '--socks', help='Get socks', default=False)
+    parser.add_argument('-g', '--glasses', help='Get glasses', default=False)
+    megroup.add_argument('-p', '--paintjob', help='Get paintjob', default=False)
+    args = parser.parse_args()
+
+    total_data = {}
+    if args.nation:
+        total_data['country_code'] = cc.convert(names=args.nation, to='ISOnumeric')
+    if args.jersey:
+        total_data.update(get_item("./JERSEYS/JERSEY", args.jersey, 'ride_jersey'))
+    if args.bike:
+        total_data.update(get_item("./BIKEFRAMES/BIKEFRAME", args.bike, 'bike_frame'))
+    if args.wheels:
+        total_data.update(get_item("./BIKEFRONTWHEELS/BIKEFRONTWHEEL", args.wheels, 'bike_wheel_front'))
+        total_data.update(get_item("./BIKEREARWHEELS/BIKEREARWHEEL", args.wheels, 'bike_wheel_rear'))
+    if args.helmet:
+        total_data.update(get_item("./HEADGEARS/HEADGEAR", args.helmet, 'ride_helmet_type'))
+    if args.socks:
+        total_data.update(get_item("./SOCKS/SOCK", args.socks, 'ride_socks_type'))
+    if args.shoes:
+        total_data.update(get_item("./BIKESHOES/BIKESHOE", args.shoes, 'ride_shoes_type'))
+    if args.glasses:
+        total_data.update(get_item("./GLASSES/GLASS", args.glasses, 'glass_type'))
+    if args.paintjob:
+        total_data.update(get_item("./PAINTJOBS/PAINTJOB", args.paintjob, 'bike_frame_colour'))
+        total_data['bike_frame_colour'] <<= 32
+        total_data.update(get_item("./BIKEFRAMES/BIKEFRAME", total_data['bike_frame_colour_name'].split('-')[0], 'bike_frame'))
+
+    data = json.dumps(total_data, indent=2)
+    print(data)
+    open('ghost_profile.txt', 'w').write(data)
+
+if __name__ == '__main__':
+    try:
+        main(sys.argv)
+    except KeyboardInterrupt:
+        pass
+    except SystemExit as se:
+        print("ERROR:", se)
+

scripts/get_pro_names.py

@@ -11,7 +11,6 @@
 import json
 import country_converter as coco
 import argparse
-import getpass
 import os
 import sys
 import xml.etree.ElementTree as ET
@@ -589,7 +588,7 @@ def get_pros(url, male, get_jersey, get_equipment, team_abbrv):
                             if 'bike_signature' in team:
                                 tmp['bike_frame'] = team['bike_signature']
                             if 'bike_frame_colour_signature' in team:
-                                tmp['bike_frame_colour'] = (team['bike_frame_colour_signature'] << 32)
+                                tmp['bike_frame_colour'] = team['bike_frame_colour_signature'] << 32
                             if 'front_wheel_signature' in team:
                                 tmp['bike_wheel_front'] = team['front_wheel_signature']
                             if 'rear_wheel_signature' in team:
@@ -601,7 +600,10 @@ def get_pros(url, male, get_jersey, get_equipment, team_abbrv):
 
     return data
 
-tree = ET.fromstring(urllib.request.urlopen('http://cdn.zwift.com/gameassets/GameDictionary.xml').read())
+gd_file = 'GameDictionary.xml'
+if not os.path.isfile(gd_file):
+    open(gd_file, 'wb').write(urllib.request.urlopen('http://cdn.zwift.com/gameassets/%s' % gd_file).read())
+tree = ET.parse(gd_file)
 jerseys = {}
 for x in tree.findall("./JERSEYS/JERSEY"):
     jerseys[x.get('name')] = int(x.get('signature'))