From e6425b55cbafc69f2db0cdbe40d701078956dba8 Mon Sep 17 00:00:00 2001 From: Ozan Goktan Date: Fri, 15 Dec 2023 10:23:36 +0100 Subject: [PATCH] add equals and hash to TimeIntervalConfig --- CHANGELOG.md | 6 +++ cognite/extractorutils/__init__.py | 2 +- .../extractorutils/configtools/elements.py | 50 ++++++++++++++++--- pyproject.toml | 2 +- 4 files changed, 50 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ffc44cf..c321ebe1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,12 @@ Changes are grouped as follows - `Security` in case of vulnerabilities. +## [6.1.1] + +### Fixed + + * Correctly handle equality comparison of `TimeIntervalConfig` objects. + ## [6.1.0] ### Added diff --git a/cognite/extractorutils/__init__.py b/cognite/extractorutils/__init__.py index 7c10f06f..6020eb50 100644 --- a/cognite/extractorutils/__init__.py +++ b/cognite/extractorutils/__init__.py @@ -16,5 +16,5 @@ Cognite extractor utils is a Python package that simplifies the development of new extractors. """ -__version__ = "6.1.0" +__version__ = "6.1.1" from .base import Extractor diff --git a/cognite/extractorutils/configtools/elements.py b/cognite/extractorutils/configtools/elements.py index ec25be06..51232d63 100644 --- a/cognite/extractorutils/configtools/elements.py +++ b/cognite/extractorutils/configtools/elements.py @@ -27,12 +27,25 @@ from prometheus_client import REGISTRY, start_http_server from cognite.client import ClientConfig, CogniteClient -from cognite.client.credentials import CredentialProvider, OAuthClientCertificate, OAuthClientCredentials +from cognite.client.credentials import ( + CredentialProvider, + OAuthClientCertificate, + OAuthClientCredentials, +) from cognite.client.data_classes import Asset, DataSet, ExtractionPipeline from cognite.extractorutils.configtools._util import _load_certificate_data from cognite.extractorutils.exceptions import InvalidConfigError -from cognite.extractorutils.metrics import AbstractMetricsPusher, CognitePusher, PrometheusPusher -from cognite.extractorutils.statestore import AbstractStateStore, LocalStateStore, NoStateStore, RawStateStore +from cognite.extractorutils.metrics import ( + AbstractMetricsPusher, + CognitePusher, + PrometheusPusher, +) +from cognite.extractorutils.statestore import ( + AbstractStateStore, + LocalStateStore, + NoStateStore, + RawStateStore, +) from cognite.extractorutils.util import EitherId _logger = logging.getLogger(__name__) @@ -93,6 +106,14 @@ class TimeIntervalConfig(yaml.YAMLObject): def __init__(self, expression: str) -> None: self._interval, self._expression = TimeIntervalConfig._parse_expression(expression) + def __eq__(self, other: object) -> bool: + if not isinstance(other, TimeIntervalConfig): + return NotImplemented + return self._interval == other._interval + + def __hash__(self) -> int: + return hash(self._interval) + @classmethod def _parse_expression(cls, expression: str) -> Tuple[int, str]: # First, try to parse pure number and assume seconds (for backwards compatibility) @@ -170,7 +191,10 @@ def _parse_expression(cls, expression: str) -> Tuple[int, str]: expression_lower = expression.lower() for size in sizes: if expression_lower.endswith(size): - return int(float(expression_lower.replace(size, "")) * sizes[size]), expression + return ( + int(float(expression_lower.replace(size, "")) * sizes[size]), + expression, + ) else: raise InvalidConfigError(f"Invalid unit for file size: {expression}. Valid units: {sizes.keys()}") @@ -241,7 +265,10 @@ class CogniteConfig: host: str = "https://api.cognitedata.com" def get_cognite_client( - self, client_name: str, token_custom_args: Optional[Dict[str, str]] = None, use_experimental_sdk: bool = False + self, + client_name: str, + token_custom_args: Optional[Dict[str, str]] = None, + use_experimental_sdk: bool = False, ) -> CogniteClient: from cognite.client.config import global_config @@ -264,7 +291,8 @@ def get_cognite_client( else: raise InvalidConfigError("Either authority-url or tenant is required for certificate authentication") (thumprint, key) = _load_certificate_data( - self.idp_authentication.certificate.path, self.idp_authentication.certificate.password + self.idp_authentication.certificate.path, + self.idp_authentication.certificate.password, ) credential_provider = OAuthClientCertificate( authority_url=authority_url, @@ -479,7 +507,10 @@ def start_pushers(self, cdf_client: CogniteClient, cancellation_token: Event = E asset = None if self.cognite.asset_name is not None and self.cognite.asset_external_id: - asset = Asset(name=self.cognite.asset_name, external_id=self.cognite.asset_external_id) + asset = Asset( + name=self.cognite.asset_name, + external_id=self.cognite.asset_external_id, + ) pusher = CognitePusher( cdf_client=cdf_client, @@ -586,7 +617,10 @@ def create_state_store( ) if self.local: - return LocalStateStore(file_path=self.local.path, save_interval=self.local.save_interval.seconds) + return LocalStateStore( + file_path=self.local.path, + save_interval=self.local.save_interval.seconds, + ) if default_to_local: return LocalStateStore(file_path="states.json") diff --git a/pyproject.toml b/pyproject.toml index 70a759bc..dd0e7e11 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "cognite-extractor-utils" -version = "6.1.0" +version = "6.1.1" description = "Utilities for easier development of extractors for CDF" authors = ["Mathias Lohne "] license = "Apache-2.0"