refactor: specifications_url (like property)

feat: get access to inner support of data sources test: refactor by new format
QuantConnect · Apr 29, 2024 · 14c4638 · 14c4638
1 parent a902ac6
commit 14c4638
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 28 deletions.
diff --git a/lean/commands/data/download.py b/lean/commands/data/download.py
@@ -415,7 +415,9 @@ def _get_available_datasets(organization: QCFullOrganization) -> List[Dataset]:
 def _get_historical_data_provider() -> str:
     return container.logger.prompt_list("Select a historical data provider", [Option(id=data_downloader.get_name(), label=data_downloader.get_name()) for data_downloader in cli_data_downloaders])
 
-def _get_param_from_config(data_provider_config_json: Dict[str, Any], default_param: List[str], key_config_data: str) -> List[str]:
+
+def _get_download_specification_from_config(data_provider_config_json: Dict[str, Any], default_param: List[str],
+                                            key_config_data: str) -> List[str]:
     """
     Get parameter from data provider config JSON or return default parameters.
 
@@ -427,10 +429,12 @@ def _get_param_from_config(data_provider_config_json: Dict[str, Any], default_pa
     Returns:
     - List[str]: List of parameters.
     """
-    if data_provider_config_json is None:
-        return default_param
 
-    return data_provider_config_json.get(key_config_data, default_param)
+    if data_provider_config_json and "module-specification" in data_provider_config_json:
+        if "download" in data_provider_config_json["module-specification"]:
+            return data_provider_config_json["module-specification"]["download"].get(key_config_data, default_param)
+
+    return default_param
 
 
 def _get_user_input_or_prompt(user_input_data: str, available_input_data: List[str], data_provider_name: str,
@@ -593,22 +597,25 @@ def download(ctx: Context,
         all_data_files = _get_data_files(organization, products)
         container.data_downloader.download_files(all_data_files, overwrite, organization.id)
     else:
-        data_downloader_provider = next(data_downloader for data_downloader in cli_data_downloaders if data_downloader.get_name() == data_provider_historical)
+        data_downloader_provider = next(data_downloader for data_downloader in cli_data_downloaders
+                                        if data_downloader.get_name() == data_provider_historical)
 
         data_provider_config_json = None
-        if data_downloader_provider._specifications_url is not None:
-            data_provider_config_json = container.api_client.data.download_public_file_json(data_downloader_provider._specifications_url)
-
-        data_provider_support_security_types = _get_param_from_config(data_provider_config_json,
-                                                                      QCSecurityType.get_all_members(),
-                                                                      "data-supported")
-        data_provider_support_data_types = _get_param_from_config(data_provider_config_json,
-                                                                  QCDataType.get_all_members(),
-                                                                  "data-types")
-        data_provider_support_resolutions = _get_param_from_config(data_provider_config_json,
-                                                                   QCResolution.get_all_members(),
-                                                                   "data-resolutions")
-        data_provider_support_markets = _get_param_from_config(data_provider_config_json, [market], "data-markets")
+        if data_downloader_provider.specifications_url is not None:
+            data_provider_config_json = container.api_client.data.download_public_file_json(
+                data_downloader_provider.specifications_url)
+
+        data_provider_support_security_types = _get_download_specification_from_config(data_provider_config_json,
+                                                                                       QCSecurityType.get_all_members(),
+                                                                                       "data-supported")
+        data_provider_support_data_types = _get_download_specification_from_config(data_provider_config_json,
+                                                                                   QCDataType.get_all_members(),
+                                                                                   "data-types")
+        data_provider_support_resolutions = _get_download_specification_from_config(data_provider_config_json,
+                                                                                    QCResolution.get_all_members(),
+                                                                                    "data-resolutions")
+        data_provider_support_markets = _get_download_specification_from_config(data_provider_config_json,
+                                                                                [market], "data-markets")
 
         security_type = _get_user_input_or_prompt(security_type, data_provider_support_security_types,
                                                   data_provider_historical, "Select a Ticker's security type")

diff --git a/lean/models/json_module.py b/lean/models/json_module.py
@@ -262,6 +262,10 @@ def _save_property(self, settings: Dict[str, Any]):
         from lean.container import container
         container.lean_config_manager.set_properties(settings)
 
+    @property
+    def specifications_url(self):
+        return self._specifications_url
+
 
 class LiveInitialStateInput(str, Enum):
     Required = "required"

diff --git a/tests/commands/data/test_download.py b/tests/commands/data/test_download.py
@@ -51,19 +51,25 @@ def _get_data_provider_config(is_crypto_configs: bool = False) -> Dict[str, Any]
 
     if is_crypto_configs:
         return {
-            "data-types": ["Trade", "Quote"],
-            "data-resolutions": ["Minute", "Hour", "Daily"],
-            "data-supported": ["Crypto", "CryptoFuture"],
-            "data-markets": ["Binance", "Kraken"]
+            "module-specification": {
+                "download": {
+                    "data-types": ["Trade", "Quote"],
+                    "data-resolutions": ["Minute", "Hour", "Daily"],
+                    "data-supported": ["Crypto", "CryptoFuture"],
+                    "data-markets": ["Binance", "Kraken"]
+                }
+            }
         }
 
     data_provider_config_file_json: Dict[str, Any] = {
-        "data-types": ["Trade", "Quote"],  # Supported data types: Trade and Quote
-        "data-resolutions": ["Second", "Minute", "Hour", "Daily"],
-        # Supported data resolutions: Second, Minute, Hour, Daily
-        "data-supported": ["Equity", "Option", "Index", "IndexOption"],
-        # Supported asset classes: Equity, Equity Options, Indexes, Index Options
-        "data-markets": ["NYSE", "USA"]
+        "module-specification": {
+            "download": {
+                "data-types": ["Trade", "Quote"],
+                "data-resolutions": ["Second", "Minute", "Hour", "Daily"],
+                "data-supported": ["Equity", "Option", "Index", "IndexOption"],
+                "data-markets": ["NYSE", "USA"]
+            }
+        }
     }
 
     return data_provider_config_file_json