draft

music_kraken/download/__init__.py

@@ -8,8 +8,8 @@ from copy import copy
 from dataclasses import dataclass, field
 from pathlib import Path
 from string import Formatter
-from typing import (TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Type,
-                    TypedDict, Union)
+from typing import (TYPE_CHECKING, Any, Callable, Dict, Generator, List,
+                    Optional, Set, Tuple, Type, TypedDict, Union)
 
 import requests
 from bs4 import BeautifulSoup
@@ -23,7 +23,7 @@ from ..utils import BColors, output, trace
 from ..utils.config import main_settings, youtube_settings
 from ..utils.enums import ALL_SOURCE_TYPES, SourceType
 from ..utils.enums.album import AlbumType
-from ..utils.exception import MKMissingNameException
+from ..utils.exception import MKComposeException, MKMissingNameException
 from ..utils.exception.download import UrlNotFoundException
 from ..utils.path_manager import LOCATIONS
 from ..utils.shared import DEBUG_PAGES
@@ -75,7 +75,7 @@ class Downloader:
             self.scan_for_pages(**kwargs)
 
     def register_page(self, page_type: Type[Page], **kwargs):
-        if page_type in _registered_pages:
+        if page_type in self._registered_pages:
             return
 
         self._registered_pages[page_type].add(page_type(
@@ -343,13 +343,15 @@ class Page:
         return super().__new__(cls)
 
     def __init__(self, download_options: DownloadOptions = None, fetch_options: FetchOptions = None, **kwargs):
-        self.SOURCE_TYPE.register_page(self)
+        if self.SOURCE_TYPE is not None:
+            self.SOURCE_TYPE.register_page(self)
         
         self.download_options: DownloadOptions = download_options or DownloadOptions()
         self.fetch_options: FetchOptions = fetch_options or FetchOptions()
 
     def __del__(self):
-        self.SOURCE_TYPE.deregister_page()
+        if self.SOURCE_TYPE is not None:
+            self.SOURCE_TYPE.deregister_page()
 
     def _search_regex(self, pattern, string, default=None, fatal=True, flags=0, group=None):
         """
@@ -451,27 +453,106 @@ class Option:
     This could represent a data object, a string or a page.
     """
 
-    def __init__(self, value: Any, text: Optional[str] = None, keys: Set[str] = None, hidden: bool = False):
+    def __init__(
+        self, 
+        value: Any, 
+        text: Optional[str] = None, 
+        keys: List[Any] = None, 
+        hidden: bool = False, 
+        parse_key: Callable[[Any], Any] = lambda x: x,
+    ):
+        self._parse_key: Callable[[Any], Any] = parse_key
+        
         self.value = value
         self.text = text or str(value)
         self.hidden = hidden
 
-        self.keys = keys or set()
-        self.keys.add(self.text)
+        self._raw_keys = set(keys or [])
+        self._raw_keys.add(self.text)
+        self.keys = set(self.parse_key(key) for key in self._raw_keys)
+    
+    def register_key(self, key: Any):
+        self._raw_keys.add(key)
+        self.keys.add(self._parse_key(key))
+
+    @property
+    def parse_key(self) -> Callable[[Any], Any]:
+        return self._parse_key
+
+    @parse_key.setter
+    def parse_key(self, value: Callable[[Any], Any]):
+        self._parse_key = value
+        self.keys = set(self._parse_key(key) for key in self._raw_keys)
+
+    def __str__(self):
+        return self.text
 
 
-class SelectOption:
-    def __init__(self, options: List[Option] = None):
+class Select:
+    def __init__(
+        self, 
+        options: List[Option] = None, 
+        option_factory: Callable[[Any], Option] = None,
+        raw_options: List[Any] = None,
+        parse_option_key: Callable[[Any], Any] = lambda x: x,
+        ask_for_creating_option: Callable[[Option], bool] = lambda x: True,
+        **kwargs
+    ):
+        self._parse_option_key: Callable[[Any], Any] = parse_option_key
+        self._ask_for_creating_option: Callable[[Option], bool] = ask_for_creating_option
+
         self._key_to_option: Dict[Any, Option] = dict()
-        self._options: List[Option] = options
+        self._options: List[Option] = []
 
-        self.extend(options or [])
+        options = options or []
+        self.option_factory: Optional[Callable[[Any], Option]] = option_factory
+        if self.can_create_options:
+            for raw_option in raw_options or []:
+                self.append(self.option_factory(raw_option))
+        elif raw_options is not None:
+            raise MKComposeException("Cannot create options without a factory.")
+
+        self.extend(options)
+
+    @property
+    def can_create_options(self) -> bool:
+        return self.option_factory is not None
 
     def append(self, option: Option):
+        option.parse_key = self._parse_option_key
         self._options.append(option)
         for key in option.keys:
             self._key_to_option[key] = option
 
     def extend(self, options: List[Option]):
         for option in options:
-            self.append(option)
+            self.append(option)
+
+    def __iter__(self) -> Generator[Option, None, None]:
+        for option in self._options:
+            if option.hidden:
+                continue
+            
+            yield option
+
+    def __contains__(self, key: Any) -> bool:
+        return key in self._key_to_option
+
+    def __getitem__(self, key: Any) -> Option:
+        return self._key_to_option[key]
+
+    def create_option(self, key: Any, **kwargs) -> Option:
+        if not self.can_create_options:
+            raise MKComposeException("Cannot create options without a factory.")
+
+        option = self.option_factory(key, **kwargs)
+        self.append(option)
+        return option
+
+    def choose(self, key: Any) -> Optional[Option]:
+        if key not in self:
+            if self.can_create_options and self._ask_for_creating_option(key):
+                return self.create_option(key)
+            return None
+
+        return self[key]

music_kraken/download/results.py

@@ -1,8 +1,12 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import Dict, Generator, List, Tuple, Type, Union
+from typing import TYPE_CHECKING, Dict, Generator, List, Tuple, Type, Union
 
 from ..objects import DatabaseObject
-from . import Page
+
+if TYPE_CHECKING:
+    from . import Page
 
 
 @dataclass