Make State a first class object

Author: judahrand

setup.py

@@ -15,7 +15,8 @@
           'simplejson==3.11.1',
           'python-dateutil>=2.6.0',
           'backoff==1.8.0',
-	  'ciso8601',
+	      'ciso8601',
+          'typing-extensions'
       ],
       extras_require={
           'dev': [

singer/__init__.py

@@ -61,17 +61,7 @@
 )
 from singer.schema import Schema
 
-from singer.bookmarks import (
-    write_bookmark,
-    get_bookmark,
-    clear_bookmark,
-    reset_stream,
-    set_offset,
-    clear_offset,
-    get_offset,
-    set_currently_syncing,
-    get_currently_syncing,
-)
+from singer.bookmarks import State
 
 if __name__ == "__main__":
     import doctest

singer/bookmarks.py

@@ -1,46 +1,96 @@
-def ensure_bookmark_path(state, path):
-    submap = state
-    for path_component in path:
-        if submap.get(path_component) is None:
-            submap[path_component] = {}
-
-        submap = submap[path_component]
-    return state
-
-def write_bookmark(state, tap_stream_id, key, val):
-    state = ensure_bookmark_path(state, ['bookmarks', tap_stream_id])
-    state['bookmarks'][tap_stream_id][key] = val
-    return state
-
-def clear_bookmark(state, tap_stream_id, key):
-    state = ensure_bookmark_path(state, ['bookmarks', tap_stream_id])
-    state['bookmarks'][tap_stream_id].pop(key, None)
-    return state
-
-def reset_stream(state, tap_stream_id):
-    state = ensure_bookmark_path(state, ['bookmarks', tap_stream_id])
-    state['bookmarks'][tap_stream_id] = {}
-    return state
-
-def get_bookmark(state, tap_stream_id, key, default=None):
-    return state.get('bookmarks', {}).get(tap_stream_id, {}).get(key, default)
-
-def set_offset(state, tap_stream_id, offset_key, offset_value):
-    state = ensure_bookmark_path(state, ['bookmarks', tap_stream_id, "offset", offset_key])
-    state['bookmarks'][tap_stream_id]["offset"][offset_key] = offset_value
-    return state
-
-def clear_offset(state, tap_stream_id):
-    state = ensure_bookmark_path(state, ['bookmarks', tap_stream_id, "offset"])
-    state['bookmarks'][tap_stream_id]["offset"] = {}
-    return state
-
-def get_offset(state, tap_stream_id, default=None):
-    return state.get('bookmarks', {}).get(tap_stream_id, {}).get("offset", default)
-
-def set_currently_syncing(state, tap_stream_id):
-    state['currently_syncing'] = tap_stream_id
-    return state
-
-def get_currently_syncing(state, default=None):
-    return state.get('currently_syncing', default)
+import json
+import sys
+from typing import Any, Dict, Optional, Sequence, Union
+from .logger import get_logger
+
+
+LOGGER = get_logger()
+
+def write_state(state):
+    json.dump(state.to_dict(), sys.stdout, indent=2)
+
+class State:
+    def __init__(
+        self, bookmarks: Optional[Dict] = None, currently_syncing: Optional[str] = None
+    ) -> None:
+        self._bookmarks = bookmarks or {}
+        self._currently_syncing = currently_syncing
+
+    def __str__(self) -> str:
+        return str(self.__dict__)
+
+    def __eq__(self, other: Any) -> bool:
+        return self.__dict__ == other.__dict__
+
+    @property
+    def bookmarks(self) -> Dict:
+        return self._bookmarks
+
+    @classmethod
+    def load(cls, filename: str) -> "State":
+        with open(filename) as fp:  # pylint: disable=invalid-name
+            return State.from_dict(json.load(fp))
+
+    @classmethod
+    def from_dict(cls, data: Dict) -> "State":
+        return State(
+            bookmarks=data.get("bookmarks"),
+            currently_syncing=data.get("currently_syncing"),
+        )
+
+    def to_dict(self) -> Dict:
+        state: Dict[str, Any] = {"bookmarks": self.bookmarks}
+        if self.get_currently_syncing():
+            state["currently_syncing"] = self.get_currently_syncing()
+        return state
+
+    def dump(self) -> None:
+        json.dump(self.to_dict(), sys.stdout, indent=2)
+
+    def _ensure_bookmark_path(self, path: Sequence) -> None:
+        submap = self.bookmarks
+        for path_component in path:
+            if submap.get(path_component) is None:
+                submap[path_component] = {}
+
+            submap = submap[path_component]
+
+    def write_bookmark(self, tap_stream_id: str, key: str, val: Any) -> None:
+        self._ensure_bookmark_path((tap_stream_id,))
+        self.bookmarks[tap_stream_id][key] = val
+
+    def clear_bookmark(self, tap_stream_id: str, key: str) -> None:
+        self._ensure_bookmark_path((tap_stream_id,))
+        self.bookmarks[tap_stream_id].pop(key, None)
+
+    def reset_stream(self, tap_stream_id: str) -> None:
+        self._ensure_bookmark_path((tap_stream_id,))
+        self.bookmarks[tap_stream_id] = {}
+
+    def get_bookmark(self, tap_stream_id: str, key: str, default: Any = None) -> Any:
+        return self.bookmarks.get(tap_stream_id, {}).get(key, default)
+
+    def set_offset(
+        self, tap_stream_id: str, offset_key: str, offset_value: Any
+    ) -> None:
+        self._ensure_bookmark_path((tap_stream_id, "offset", offset_key))
+        self.bookmarks[tap_stream_id]["offset"][offset_key] = offset_value
+
+    def clear_offset(self, tap_stream_id: str) -> None:
+        self._ensure_bookmark_path((tap_stream_id, "offset"))
+        self.bookmarks[tap_stream_id]["offset"] = {}
+
+    def get_offset(
+        self, tap_stream_id: str, offset_key: str, default: Any = None
+    ) -> Any:
+        return (
+            self.bookmarks.get(tap_stream_id, {})
+            .get("offset", {})
+            .get(offset_key, default)
+        )
+
+    def get_currently_syncing(self, default: Optional[str] = None) -> Optional[str]:
+        return self._currently_syncing or default
+
+    def set_currently_syncing(self, value: Union[str, None]) -> None:
+        self._currently_syncing = value

singer/catalog.py

@@ -3,7 +3,6 @@
 import sys
 
 from . import metadata as metadata_module
-from .bookmarks import get_currently_syncing
 from .logger import get_logger
 from .schema import Schema
 
@@ -132,7 +131,7 @@ def get_stream(self, tap_stream_id):
         return None
 
     def _shuffle_streams(self, state):
-        currently_syncing = get_currently_syncing(state)
+        currently_syncing = state.get_currently_syncing()
 
         if currently_syncing is None:
             return self.streams