support collections

steam-workshop-download

@@ -3,43 +3,83 @@
 from requests import Session
 from os import utime
 from os.path import exists, getmtime, getatime, getsize, abspath, join
-from sys import argv
 import argparse
+import re
+from concurrent.futures import ThreadPoolExecutor
 
 # args
 parser = argparse.ArgumentParser(description='Download items from steam workshop.')
-parser.add_argument('ids', metavar='ITEM_ID', nargs='+', type=int, help='steam workshop file ids')
+parser.add_argument('-i', '--items', metavar='ID', nargs='+', dest='item_ids', type=int, default=[], help='workshop item ids')
-parser.add_argument('-o', '--out', metavar='TRAGET_DIR', dest='out', type=abspath, default='.', help='output dir')
+parser.add_argument('-c', '--collections', metavar='ID', nargs='+', dest='collection_ids', type=int, default=[], help='workshop collection ids')
+parser.add_argument('-o', '--out', metavar='PATH', dest='out', type=abspath, default='.', help='output dir')
+parser.add_argument('-f', '--filename', metavar='PATTERN', dest='filename', type=str, default='{title}-{id}.{ext}', help='output file basename pattern: defaults to \'{title}-{id}.{ext}\', also supports {filename}')
+parser.add_argument('-p', '--parallel-downloads', metavar='NUMBER', dest='parallel_downloads', type=int, default=8, help='number of parallel downloads')
 args = parser.parse_args()
-print(f"item ids: {', '.join(str(id) for id in args.ids)}", flush=True)
+
-print(f"target dir: {args.out}", flush=True)
+if not args.item_ids and not args.collection_ids:
+    parser.error('at least one of --items or --collections must be specified')
 
 # init http session
 session = Session()
 
-# get item information
+def flash(string):
+    print(f"\r\033[K{string}...", end="", flush=True)
+
+def log(string):
+    print(f"\r\033[K{string}")
+
+# get items from collections
+if args.collection_ids:
+    flash('resolving collections')
+    response = session.post(
+        'http://api.steampowered.com/ISteamRemoteStorage/GetCollectionDetails/v1',
+        data={
+            'collectioncount' : len(args.collection_ids),
+            **{
+                f'publishedfileids[{i}]' : collection_id
+                    for i, collection_id in enumerate(args.collection_ids)
+            },
+        },
+    )
+    response.raise_for_status()
+
+    # extract item ids
+    for collection in response.json()['response']['collectiondetails']:
+        for child in collection.get('children', []):
+            if child['filetype'] == 0:
+                args.item_ids.append(int(child['publishedfileid']))
+
+# get items
+flash('resolving items')
 response = session.post(
     'http://api.steampowered.com/ISteamRemoteStorage/GetPublishedFileDetails/v1',
     data={
-        'itemcount' : len(args.ids),
+        'itemcount' : len(args.item_ids),
         **{
             f'publishedfileids[{i}]' : plugin_id
-                for i, plugin_id in enumerate(args.ids)
+                for i, plugin_id in enumerate(args.item_ids)
         },
     },
 )
 response.raise_for_status()
 
-# download items
+# create thread pool
-for item in response.json()['response']['publishedfiledetails']:
+executor = ThreadPoolExecutor(max_workers=args.parallel_downloads)
+
+# download function
+def download(item):
     # file found?
     if item['result'] != 1:
         raise ValueError(f"getting file '{item['publishedfileid']}' info failed: {item}")
 
     # get target path
-
+    filename_without_ext, ext = item['filename'].split('/')[-1].rsplit('.', 1)
-    target_path = join(args.out, f"{item['publishedfileid']}.vpk")
+    target_path = join(args.out, f"{args.filename}".format(
-    print(f"- {item['title']}: ", end='', flush=True)
+        id=item['publishedfileid'],
+        title=re.sub(r'[^a-zA-Z0-9]+', '_', item['title']).strip('_'),
+        filename=filename_without_ext,
+        ext=ext,
+    ))
 
     # skip item?
     if (
@@ -47,24 +87,36 @@ for item in response.json()['response']['publishedfiledetails']:
         int(item['time_updated']) == int(getmtime(target_path)) and # mtime matches
         int(item['file_size']) == int(getsize(target_path)) # filesize matches
     ):
-        print(f"already satisfied", flush=True)
+        log(f'skipped {target_path}')
-        continue
+        return
 
     # download item
-    print(f"downloading", end='', flush=True)
     response = session.get(item['file_url'], stream=True)
     response.raise_for_status()
 
     # one chunk at a time
     with open(target_path, 'wb') as file:
-        for chunk in response.iter_content(chunk_size=65_536):
+        written = 0
-            print('.', end='', flush=True)
+        for chunk in response.iter_content(chunk_size=1_048_576):
             if chunk:
+                # print every 8MB
+                written += len(chunk)
+                if written // (8 * 1024 * 1024) > 0:
+                    flash(f'downloading {target_path}')
+                    written = 0
+
                 file.write(chunk)
 
     # update modify time
-    print(' done', flush=True)
+    log(f'finished {target_path}')
     utime(target_path, (getatime(target_path), item['time_updated']))
 
+# download items
+for item in response.json()['response']['publishedfiledetails']:
+    executor.submit(download, item)
+
+# wait for all downloads to finish
+executor.shutdown(wait=True)
+
 # close session
 session.close()