From abfa74fbb554397544273eb73f965234944444e6 Mon Sep 17 00:00:00 2001 From: Alexandre Catarino Date: Fri, 3 Nov 2023 22:26:17 +0000 Subject: [PATCH] LEAN CLI will verify whether the Organization has a Security Master requirement on the Dataset page. For example, US Futures requires US Futures Security Master. If the Org. doesn't have it, LEAN CLI directs the user to the US Futures Security Master pricing page. (#378) Tested locally with Orgs. with and without Futures Security Master, and setting and not setting the requirements ```json "requirements": { "137": "quantconnect-us-futures-security-master" }, ``` Closes QuantConnect#73 --- lean/commands/data/download.py | 23 +++++++++++++---------- lean/constants.py | 3 --- lean/models/api.py | 13 ++++--------- lean/models/data.py | 2 +- tests/commands/data/test_download.py | 24 ++++++++++++++++++++---- 5 files changed, 38 insertions(+), 27 deletions(-) diff --git a/lean/commands/data/download.py b/lean/commands/data/download.py index 267f0879..59209053 100644 --- a/lean/commands/data/download.py +++ b/lean/commands/data/download.py @@ -153,10 +153,10 @@ def _display_products(organization: QCFullOrganization, products: List[Product]) logger.info(f"Organization balance: {organization.credit.balance:,.0f} QCC") -def _get_security_master_warn() -> str: +def _get_security_master_warn(url: str) -> str: return "\n".join([f"Your organization does not have an active Security Master subscription. Override the Security Master precautions will likely" f" result in inaccurate and misleading backtest results. Use this override flag at your own risk.", - f"You can add the subscription at https://www.quantconnect.com/datasets/quantconnect-security-master/pricing" + f"You can add the subscription at https://www.quantconnect.com/datasets/{url}/pricing" ]) @@ -196,15 +196,16 @@ def _select_products_interactive(organization: QCFullOrganization, datasets: Lis dataset: Dataset = logger.prompt_list("Select a dataset", [Option(id=d, label=d.name) for d in available_datasets]) - if dataset.requires_security_master and not organization.has_security_master_subscription(): + for id, url in dataset.requirements.items(): + if organization.has_security_master_subscription(id): + continue if not force: logger.warn("\n".join([ f"Your organization needs to have an active Security Master subscription to download data from the '{dataset.name}' dataset", - f"You can add the subscription at https://www.quantconnect.com/datasets/quantconnect-security-master/pricing" + f"You can add the subscription at https://www.quantconnect.com/datasets/{url}/pricing" ])) - continue else: - logger.warn(_get_security_master_warn()) + logger.warn(_get_security_master_warn(url)) option_results = OrderedDict() for dataset_option in dataset.options: @@ -328,14 +329,16 @@ def _select_products_non_interactive(organization: QCFullOrganization, if dataset is None: raise RuntimeError(f"There is no dataset named '{ctx.params['dataset']}'") - if dataset.requires_security_master and not organization.has_security_master_subscription(): + for id, url in dataset.requirements.items(): + if organization.has_security_master_subscription(id): + continue if not force: raise RuntimeError("\n".join([ f"Your organization needs to have an active Security Master subscription to download data from the '{dataset.name}' dataset", - f"You can add the subscription at https://www.quantconnect.com/datasets/quantconnect-security-master/pricing" + f"You can add the subscription at https://www.quantconnect.com/datasets/{url}/pricing" ])) else: - container.logger.warn(_get_security_master_warn()) + container.logger.warn(_get_security_master_warn(url)) option_results = OrderedDict() invalid_options = [] @@ -401,7 +404,7 @@ def _get_available_datasets(organization: QCFullOrganization) -> List[Dataset]: categories=[tag.name.strip() for tag in cloud_dataset.tags], options=datasource["options"], paths=datasource["paths"], - requires_security_master=datasource["requiresSecurityMaster"])) + requirements=datasource.get("requirements", {}))) return available_datasets diff --git a/lean/constants.py b/lean/constants.py index 7490c229..121f9a71 100644 --- a/lean/constants.py +++ b/lean/constants.py @@ -91,9 +91,6 @@ # The interval in hours at which the CLI checks for new announcements UPDATE_CHECK_INTERVAL_ANNOUNCEMENTS = 24 -# The product id of the Equity Security Master subscription -EQUITY_SECURITY_MASTER_PRODUCT_ID = 37 - # The product id of the Terminal Link module TERMINAL_LINK_PRODUCT_ID = 44 diff --git a/lean/models/api.py b/lean/models/api.py index b77f741a..2662250e 100644 --- a/lean/models/api.py +++ b/lean/models/api.py @@ -15,7 +15,6 @@ from enum import Enum from typing import Any, Dict, List, Optional, Union -from lean.constants import EQUITY_SECURITY_MASTER_PRODUCT_ID from lean.models.pydantic import WrappedBaseModel, validator @@ -403,22 +402,18 @@ class QCFullOrganization(WrappedBaseModel): data: QCOrganizationData members: List[QCOrganizationMember] - def has_security_master_subscription(self) -> bool: - """Returns whether this organization has a Security Master subscription. + def has_security_master_subscription(self, id: int) -> bool: + """Returns whether this organization has the Security Master subscription of a given Id + :param id: the Id of the Security Master Subscription :return: True if the organization has a Security Master subscription, False if not """ - # TODO: This sort of hardcoded product ID checking is not sufficient when we consider - # multiple 'Security Master' products. Especially since they will be specific to certain datasources. - # For now, simple checks for an equity "Security Master" subscription - # Issue created here: https://github.com/QuantConnect/lean-cli/issues/73 - data_products_product = next((x for x in self.products if x.name == "Data"), None) if data_products_product is None: return False - return any(x.productId in {EQUITY_SECURITY_MASTER_PRODUCT_ID} for x in data_products_product.items) + return any(x.productId == id for x in data_products_product.items) class QCMinimalOrganization(WrappedBaseModel): diff --git a/lean/models/data.py b/lean/models/data.py index 5e98433d..19d4dcec 100644 --- a/lean/models/data.py +++ b/lean/models/data.py @@ -283,7 +283,7 @@ class Dataset(WrappedBaseModel): categories: List[str] options: List[DatasetOption] paths: List[DatasetPath] - requires_security_master: bool + requirements: Dict[int, str] @validator("options", pre=True) def parse_options(cls, values: List[Any]) -> List[Any]: diff --git a/tests/commands/data/test_download.py b/tests/commands/data/test_download.py index 4d5aa8ce..be34f851 100644 --- a/tests/commands/data/test_download.py +++ b/tests/commands/data/test_download.py @@ -46,14 +46,30 @@ def test_interactive_bulk_select(): categories=["testData"], options=datasource["options"], paths=datasource["paths"], - requires_security_master=datasource["requiresSecurityMaster"])] + requirements=datasource.get("requirements", {}))] products = _select_products_interactive(organization, testSets) # No assertion, since interactive has multiple results +def test_dataset_requirements(): + organization = create_api_organization() + datasource = json.loads(bulk_datasource) + testSet = Dataset(name="testSet", + vendor="testVendor", + categories=["testData"], + options=datasource["options"], + paths=datasource["paths"], + requirements=datasource.get("requirements", {})) + + for id, name in testSet.requirements.items(): + assert not organization.has_security_master_subscription(id) + assert id==39 + bulk_datasource=""" { - "requiresSecurityMaster": true, + "requirements": { + "39": "quantconnect-us-equity-security-master" + }, "options": [ { "type": "select", @@ -262,12 +278,12 @@ def test_filter_pending_datasets() -> None: str(test_datasets[0].id): { 'options': [], 'paths': [], - 'requiresSecurityMaster': True, + 'requirements': {}, }, str(test_datasets[1].id): { 'options': [], 'paths': [], - 'requiresSecurityMaster': True, + 'requirements': {}, } } container.api_client.data.get_info = MagicMock(return_value=QCDataInformation(datasources=datasources, prices=[],