music_kraken/cli/main_downloader.py

@@ -1,7 +1,7 @@
 import random
 import re
 from pathlib import Path
-from typing import Dict, Generator, List, Set, Type
+from typing import Dict, Generator, List, Set, Type, Union
 
 from .. import console
 from ..download import Downloader, Page, components
@@ -112,8 +112,8 @@ class CliDownloader:
 
         print()
 
-    def set_current_options(self, current_options: Generator[DatabaseObject, None, None]):
-        current_options = components.DataObjectSelect(current_options)
+    def set_current_options(self, current_options: Union[Generator[DatabaseObject, None, None], components.Select]):
+        current_options = current_options if isinstance(current_options, components.Select) else components.DataObjectSelect(current_options)
 
         if main_settings["result_history"]:
             self._result_history.append(current_options)
@@ -221,12 +221,14 @@ class CliDownloader:
         self.set_current_options(self.downloader.search(parsed_query))
         self.print_current_options()
 
-    def goto(self, data_object: DatabaseObject):
+    def goto(self, data_object: Union[DatabaseObject, components.Select]):
         page: Type[Page]
 
-        self.downloader.fetch_details(data_object, stop_at_level=1)
-
-        self.set_current_options(GoToResults(data_object.options, max_items_per_page=self.max_displayed_options))
+        if isinstance(data_object, components.Select):
+            self.set_current_options(data_object)
+        else:
+            self.downloader.fetch_details(data_object, stop_at_level=1)
+            self.set_current_options(data_object.options)
 
         self.print_current_options()
 
@@ -293,24 +295,15 @@ class CliDownloader:
 
                 indices = []
                 for possible_index in q.split(","):
-                    possible_index = possible_index.strip()
                     if possible_index == "":
                         continue
+
+                    if possible_index not in self.current_results:
+                        raise MKInvalidInputException(message=f"The index \"{possible_index}\" is not in the current options.")
                     
-                    i = 0
-                    try:
-                        i = int(possible_index)
-                    except ValueError:
-                        raise MKInvalidInputException(message=f"The index \"{possible_index}\" is not a number.")
+                    yield self.current_results[possible_index]
 
-                    if i < 0 or i >= len(self.current_results):
-                        raise MKInvalidInputException(message=f"The index \"{i}\" is not within the bounds of 0-{len(self.current_results) - 1}.")
-                    
-                    indices.append(i)
-
-                return [self.current_results[i] for i in indices]
-
-            selected_objects = get_selected_objects(query)
+            selected_objects = list(get_selected_objects(query))
 
             if do_merge:
                 old_selected_objects = selected_objects
@@ -337,7 +330,7 @@ class CliDownloader:
             if len(selected_objects) != 1:
                 raise MKInvalidInputException(message="You can only go to one object at a time without merging.")
 
-            self.goto(selected_objects[0])
+            self.goto(selected_objects[0].value)
             return False
         except MKInvalidInputException as e:
             output("\n" + e.message + "\n", color=BColors.FAIL)

music_kraken/download/components.py

@@ -125,7 +125,11 @@ class Select:
         return self._parse_option_key(key) in self._key_to_option
 
     def __getitem__(self, key: Any) -> Option:
-        return self._key_to_option[self._parse_option_key(key)]
+        r = self._key_to_option[self._parse_option_key(key)]
+        # check if is callable
+        if callable(r.value):
+            r.value = r.value()
+        return r
 
     def create_option(self, key: Any, **kwargs) -> Option:
         if not self.can_create_options:
@@ -186,20 +190,36 @@ class GenreSelect(StringSelect):
         super().__init__(sort=True, raw_options=(genre.name for genre in filter(self.is_valid_genre, main_settings["music_directory"].iterdir())))
 
 
+
+class SourceTypeToOption(dict):
+    def __init__(self, callback):
+        super().__init__()
+        
+        self.callback = callback
+
+    def __missing__(self, key):
+        self[key] = self.callback(key)
+        return self[key]
+
+
 class DataObjectSelect(Select):
     def __init__(self, data_objects: Generator[DataObject]):
         self._source_type_to_data_objects: Dict[SourceType, List[Option]] = defaultdict(list)
+        self._source_type_to_option: Dict[SourceType, Option] = SourceTypeToOption(self.option_from_source_type)
 
         self._data_object_index: int = 0
         self._source_type_index: int = 0
 
-        super().__init__()
+        super().__init__(
+            parse_option_key=lambda x: unify(str(x)),
+        )
 
         self.extend(data_objects)
 
     def option_from_data_object(self, data_object: DataObject) -> Option:
         index = self._data_object_index
         self._data_object_index += 1
+
         return Option(
             value=data_object,
             keys=[index, data_object.option_string, data_object.title_string],
@@ -210,12 +230,16 @@ class DataObjectSelect(Select):
         index = ALPHABET[self._source_type_index % len(ALPHABET)]
         self._source_type_index += 1
 
-        return Option(
-            value=self._source_type_to_data_objects[source_type],
+        o = Option(
+            value=lambda: DataObjectSelect(self._source_type_to_data_objects[source_type]),
             keys=[index, source_type],
             text=f"{BColors.HEADER.value}({index}) --------------------------------{source_type.name:{'-'}<{21}}--------------------{BColors.ENDC.value}",
         )
 
+        super().append(o)
+
+        return o
+
     def append(self, option: Union[Option, DataObject]):
         if isinstance(option, DataObject):
             data_object = option
@@ -224,15 +248,19 @@ class DataObjectSelect(Select):
             data_object = option.value
 
         for source_type in data_object.source_collection.source_types(only_with_page=True):
-            if source_type not in self._source_type_to_data_objects:
-                st_option = self.option_from_source_type(source_type)
-                self._source_type_to_data_objects[source_type].append(st_option)
-                super().append(st_option)
-
             self._source_type_to_data_objects[source_type].append(option)
 
         super().append(option)
 
     def __iter__(self):
-        for options in self._source_type_to_data_objects.values():
-            yield from options
+        source_types = list(sorted(self._source_type_to_data_objects.keys(), key=lambda x: x.name))
+        has_limit = len(source_types) > 1
+
+        for st in source_types:
+            if has_limit:
+                yield self._source_type_to_option[st]
+
+            limit = min(15, len(self._source_type_to_data_objects[st])) if has_limit else len(self._source_type_to_data_objects[st])
+
+            for i in range(limit):
+                yield self._source_type_to_data_objects[st][i]