diff --git a/lean/constants.py b/lean/constants.py index 64e591c4..be2b4979 100644 --- a/lean/constants.py +++ b/lean/constants.py @@ -69,8 +69,11 @@ # The interval in hours at which the CLI checks for new announcements UPDATE_CHECK_INTERVAL_ANNOUNCEMENTS = 24 -# The product id of the Security Master subscription -SECURITY_MASTER_PRODUCT_ID = 37 +# The product id of the Equity Security Master subscription +EQUITY_SECURITY_MASTER_PRODUCT_ID = 37 + +# The product id of the Equity Bulk Security Master subscription +BULK_EQUITY_SECURITY_MASTER_PRODUCT_ID = 180 # 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 0b3d716d..e6dda38b 100644 --- a/lean/models/api.py +++ b/lean/models/api.py @@ -21,7 +21,7 @@ from rich.table import Table from rich.text import Text -from lean.constants import SECURITY_MASTER_PRODUCT_ID +from lean.constants import EQUITY_SECURITY_MASTER_PRODUCT_ID, BULK_EQUITY_SECURITY_MASTER_PRODUCT_ID from lean.models.pydantic import WrappedBaseModel @@ -391,11 +391,17 @@ def has_security_master_subscription(self) -> bool: :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 == SECURITY_MASTER_PRODUCT_ID for x in data_products_product.items) + return any(x.productId in {EQUITY_SECURITY_MASTER_PRODUCT_ID, BULK_EQUITY_SECURITY_MASTER_PRODUCT_ID} for x in data_products_product.items) class QCMinimalOrganization(WrappedBaseModel):