diff --git a/CHANGELOG.md b/CHANGELOG.md index d2ac52ff..7ab5e3ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,12 @@ Changes are grouped as follows - `Fixed` for any bug fixes. - `Security` in case of vulnerabilities. +## [5.5.1] + +### Added + + * Added iter method on the state store to return the keys of the local state dict + ## [5.5.0] ### Added diff --git a/cognite/extractorutils/__init__.py b/cognite/extractorutils/__init__.py index b5cecba9..6c15c82e 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__ = "5.5.0" +__version__ = "5.5.1" from .base import Extractor diff --git a/cognite/extractorutils/statestore.py b/cognite/extractorutils/statestore.py index 5a3434f5..a957666e 100644 --- a/cognite/extractorutils/statestore.py +++ b/cognite/extractorutils/statestore.py @@ -90,7 +90,7 @@ import threading from abc import ABC, abstractmethod from types import TracebackType -from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type, Union from requests.exceptions import ConnectionError @@ -309,6 +309,10 @@ def __contains__(self, external_id: str) -> bool: def __len__(self) -> int: return len(self._local_state) + def __iter__(self) -> Iterable[str]: + for key in self._local_state: + yield key + class RawStateStore(AbstractStateStore): """ diff --git a/pyproject.toml b/pyproject.toml index fde5b898..8c1f2ca4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "cognite-extractor-utils" -version = "5.5.0" +version = "5.5.1" description = "Utilities for easier development of extractors for CDF" authors = ["Mathias Lohne "] license = "Apache-2.0" diff --git a/tests/tests_unit/test_statestore.py b/tests/tests_unit/test_statestore.py index 799e681e..9093353a 100644 --- a/tests/tests_unit/test_statestore.py +++ b/tests/tests_unit/test_statestore.py @@ -116,6 +116,19 @@ def test_get_state(self): self.assertListEqual(state_store.get_state(["extId1", "extId3", "extId5"]), [(1, 5), (0, None), (None, None)]) + def test_local_state_interaction(self): + state_store = NoStateStore() + state_store._local_state = { + "extId1": {"low": 1, "high": 5}, + "extId2": {"low": None, "high": 4}, + "extId3": {"low": 0}, + "extId4": {"low": 3, "high": None}, + } + + self.assertTrue(len(state_store) == len(state_store._local_state)) + for id in state_store: + self.assertTrue(id in state_store._local_state.keys()) + @patch("cognite.client.CogniteClient") def test_upload_queue_integration(self, MockCogniteClient): state_store = NoStateStore()