Skip to content

Commit

Permalink
add equals and hash to TimeIntervalConfig
Browse files Browse the repository at this point in the history
  • Loading branch information
ozangoktan committed Dec 15, 2023
1 parent 983a049 commit e6425b5
Show file tree
Hide file tree
Showing 4 changed files with 50 additions and 10 deletions.
6 changes: 6 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion cognite/extractorutils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
50 changes: 42 additions & 8 deletions cognite/extractorutils/configtools/elements.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()}")

Expand Down Expand Up @@ -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

Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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")
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -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 <[email protected]>"]
license = "Apache-2.0"
Expand Down

0 comments on commit e6425b5

Please sign in to comment.