add_iter_method_on_state_store

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

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

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):
     """

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 <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"

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()