Merge pull request #6 from thomasasfk/tafk-feet-saber-wip

Feet saber process multiple ids at once
thomasasfk · Oct 8, 2023 · 8b4240c · 8b4240c
2 parents d0c20f7 + ee1daa0
commit 8b4240c
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 15 deletions.
diff --git a/README.md b/README.md
@@ -61,5 +61,6 @@ Generate full DD Beat Map from Feet Saber directory (WIP)
 ```bash
 .venv/Scripts/python fs2dd.py --fs-map-dir "path/to/map/folder"
 .venv/Scripts/python fs2dd.py --fs-map-id 229ed
+.venv/Scripts/python fs2dd.py --fs-map-ids 229ed,299b5
 .venv/Scripts/python fs2dd.py --fs-playlist-id 3474
 ```
diff --git a/fs2dd.py b/fs2dd.py
@@ -200,7 +200,10 @@ def create_dd_tracks_from_fs(fs_map_dir: str, prefix_dir: str = '', ost_id: int
             fs_map_dir,
         ) if f.endswith('.egg')
     ][0]
-    target_dir = f'{prefix_dir}/{fs_info.songName} - {fs_info.levelAuthorName}'
+    dir_name = f'{fs_info.songName} - {fs_info.levelAuthorName}'.replace('/', '-').replace('?', '').replace('\\', '-').replace(  # noqa
+        ':', '-',
+    ).replace('*', '-').replace("\"", '-').replace('<', '-').replace('>', '-').replace('|', '-').replace('.', '-').replace(' ', '-')  # noqa
+    target_dir = f'{prefix_dir}/{dir_name}'
     if not os.path.exists(target_dir):
         os.makedirs(target_dir)
 
@@ -295,6 +298,10 @@ def create_dd_tracks_from_fs(fs_map_dir: str, prefix_dir: str = '', ost_id: int
         '--fs-map-id', type=str,
         help='The ID to fetch the Feet Saber map zip', required=False,
     )
+    parser.add_argument(
+        '--fs-map-ids', type=str,
+        help='IDs to fetch the Feet Saber map zip, separated by comma', required=False,
+    )
     parser.add_argument(
         '--fs-playlist-id', type=str,
         help='The ID to fetch the Feet Saber playlist', required=False,
@@ -353,6 +360,49 @@ def download_and_process_track(url, target_dir, ost_id):
         print(f'Created bin/{ost_id}.zip')
         raise SystemExit(0)
 
+    if args.fs_map_ids:
+        ost_id = random_9_digit_int()
+        fs_map_ids = args.fs_map_ids.split(',')
+        target_dir = f'bin/{ost_id}'
+
+        fs_map_urls = [
+            get_feet_saber_map_from_id(fs_map_id) for fs_map_id in fs_map_ids
+        ]
+
+        def download_and_process_track(url, target_dir, ost_id):
+            fs_map_dir = download_and_extract_zip(url)
+            track = create_dd_tracks_from_fs(fs_map_dir, target_dir, ost_id)
+            return track
+
+        tracks = []
+        with ThreadPoolExecutor() as executor:
+            future_to_url = {
+                executor.submit(download_and_process_track, url, target_dir, ost_id): url for url in
+                fs_map_urls
+            }
+            for future in concurrent.futures.as_completed(future_to_url):
+                track = future.result()
+                tracks.append(track)
+
+        album_info = DDAlbumInfo(
+            OstName=f'Feet Saber - {ost_id}',
+            BeatMapIdList=sorted([track.BeatMapId for track in tracks]),
+            OstId=ost_id,
+            CoverPath='',
+            CreateTime=yyyymmdd_to_ticks(datetime.now().strftime('%Y%m%d')),
+        ).save_to_file(target_dir)
+        print(f'Created album info file: {album_info}')
+
+        print('Zipping tracks...')
+        with zipfile.ZipFile(f'bin/{ost_id}.zip', 'w', zipfile.ZIP_DEFLATED) as zipf:
+            zipdir(
+                target_dir, zipf,
+                f'Dance Dash_Data/StreamingAssets/NewDLC/{ost_id}',
+            )
+
+        print(f'Created bin/{ost_id}.zip')
+        raise SystemExit(0)
+
     if not args.fs_map_dir and not args.fs_map_id:
         print('Either --fs-map-dir or --fs-map-id must be provided.')
         raise SystemExit(1)

diff --git a/model/feetsaber.py b/model/feetsaber.py
@@ -100,7 +100,7 @@ def remove_underscore_from_keys(data):
         contributors = [
             FSContributor(**remove_underscore_from_keys(contributor)) for contributor in
             custom_data_raw['_contributors']
-        ]
+        ] if '_contributors' in custom_data_raw else []
 
         editors = {
             editor: FSEditor(version=data['version']) for editor, data in custom_data_raw['_editors'].items() if
@@ -130,10 +130,16 @@ def remove_underscore_from_keys(data):
                         difficultyLabel=beatmap['_customData'].get(
                             '_difficultyLabel',
                         ),
-                        editorOffset=beatmap['_customData']['_editorOffset'],
-                        editorOldOffset=beatmap['_customData']['_editorOldOffset'],
+                        editorOffset=beatmap['_customData'].get(
+                            '_editorOffset',
+                        ),
+                        editorOldOffset=beatmap['_customData'].get(
+                            '_editorOldOffset',
+                        ),
                         suggestions=beatmap['_customData'].get('_suggestions'),
-                        requirements=beatmap['_customData']['_requirements'],
+                        requirements=beatmap['_customData'].get(
+                            '_requirements',
+                        ),
                     ),
                 ) for beatmap in dbset['_difficultyBeatmaps']
             ]
@@ -248,9 +254,9 @@ class FSBeatMapFileObstacleCustomData:
     track: str
     interactable: bool
     fake: bool
-    position: tuple[float, float]
     scale: tuple
-    color: tuple[float, float, float, float]
+    position: tuple[float, float] | None
+    color: tuple[float, float, float, float] | None
     localRotation: tuple[float, float, float] | None = None
 
     @property
@@ -259,8 +265,9 @@ def height(self) -> float:
 
     @property
     def is_fs(self) -> bool:
-        _, y = self.position
-        return self.height == 0.1 and y == -0.25 and self.fake
+        if not self.position:
+            return False
+        return self.height == 0.1 and self.position[1] == -0.25 and self.fake
 
     @property
     def is_fs_slider(self) -> bool:
@@ -344,7 +351,7 @@ class FSBeatMapFile:
     @classmethod
     def from_json_dict(cls, json_dict: dict) -> FSBeatMapFile:
         custom_data = FSBeatMapFileCustomData(
-            time=json_dict['_customData']['_time'],
+            time=json_dict['_customData'].get('_time'),
             bookmarks=[
                 FSBeatMapFileBookmark(
                     time=bm['_time'],
@@ -371,7 +378,7 @@ def from_json_dict(cls, json_dict: dict) -> FSBeatMapFile:
                 cutDirection=note['_cutDirection'],
                 customData=FSBeatMapFileNoteCustomData(
                     position=tuple(note['_customData']['_position']),
-                ) if '_customData' in note else None,
+                ) if '_customData' in note and '_position' in note['_customData'] else None,
             ) for note in json_dict['_notes']
         ]
 
@@ -384,9 +391,11 @@ def from_json_dict(cls, json_dict: dict) -> FSBeatMapFile:
                 width=obstacle['_width'],
                 customData=FSBeatMapFileObstacleCustomData(
                     track=obstacle['_customData'].get('_track'),
-                    interactable=obstacle['_customData']['_interactable'],
-                    fake=obstacle['_customData']['_fake'],
-                    position=tuple(obstacle['_customData']['_position']),
+                    interactable=obstacle['_customData'].get('_interactable'),
+                    fake=obstacle['_customData'].get('_fake'),
+                    position=tuple(
+                        obstacle['_customData']['_position'],
+                    ) if '_position' in obstacle['_customData'] else None,
                     localRotation=tuple(obstacle['_customData']['_localRotation']) if '_localRotation' in obstacle[
                         '_customData'
                     ] else None,
@@ -395,7 +404,7 @@ def from_json_dict(cls, json_dict: dict) -> FSBeatMapFile:
                     color=tuple(
                         float(c)
                         for c in obstacle['_customData']['_color']
-                    ),
+                    ) if '_color' in obstacle['_customData'] else None,
                 ) if '_customData' in obstacle else None,
             ) for obstacle in json_dict['_obstacles']
         ]