From 46d83f2cb9ed1094c1fce444085e53ab808f1b15 Mon Sep 17 00:00:00 2001
From: lash <dev@holbrook.no>
Date: Sat, 13 Aug 2022 20:50:13 +0000
Subject: [PATCH] Add active state count, override default state name

---
 CHANGELOG           |  3 +++
 setup.cfg           |  2 +-
 shep/persist.py     |  4 ++--
 shep/state.py       | 18 ++++++++++++++----
 tests/test_state.py | 17 +++++++++++++++++
 5 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG b/CHANGELOG
index 6b7512b..4cf14c9 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,3 +1,6 @@
+- 0.2.10
+	* Add count active states method
+	* Enable complete replace of NEW state on state instantiation
 - 0.2.9
 	* Enable access to is_pure method
 	* Numeric option for elements return value
diff --git a/setup.cfg b/setup.cfg
index 87f92a6..4da45b9 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = shep
-version = 0.2.9
+version = 0.2.10
 description = Multi-state key stores using bit masks
 author = Louis Holbrook
 author_email = dev@holbrook.no
diff --git a/shep/persist.py b/shep/persist.py
index b4db0c2..1380a40 100644
--- a/shep/persist.py
+++ b/shep/persist.py
@@ -20,8 +20,8 @@ class PersistedState(State):
     :type logger: object
     """
 
-    def __init__(self, factory, bits, logger=None, verifier=None, check_alias=True, event_callback=None):
-        super(PersistedState, self).__init__(bits, logger=logger, verifier=verifier, check_alias=check_alias, event_callback=event_callback)
+    def __init__(self, factory, bits, logger=None, verifier=None, check_alias=True, event_callback=None, default_state=None):
+        super(PersistedState, self).__init__(bits, logger=logger, verifier=verifier, check_alias=check_alias, event_callback=event_callback, default_state=default_state)
         self.__store_factory = factory
         self.__stores = {}
 
diff --git a/shep/state.py b/shep/state.py
index 55481a2..4b5d153 100644
--- a/shep/state.py
+++ b/shep/state.py
@@ -30,16 +30,22 @@ class State:
 
     base_state_name = 'NEW'
 
-    def __init__(self, bits, logger=None, verifier=None, check_alias=True, event_callback=None):
+    def __init__(self, bits, logger=None, verifier=None, check_alias=True, event_callback=None, default_state=None):
         self.__initial_bits = bits
         self.__bits = bits
         self.__limit = (1 << bits) - 1
         self.__c = 0
-        setattr(self, self.base_state_name, 0)
 
-        self.__reverse = {0: getattr(self, self.base_state_name)}
-        self.__keys = {getattr(self, self.base_state_name): []}
+        if default_state == None:
+            default_state = self.base_state_name
+
+        setattr(self, default_state, 0)
+
+        self.__reverse = {0: getattr(self, default_state)}
+        self.__keys = {getattr(self, default_state): []}
         self.__keys_reverse = {}
+        if default_state != self.base_state_name:
+            self.__keys_reverse[default_state] = 0
         self.__contents = {}
         self.modified_last = {}
         self.verifier = verifier
@@ -661,3 +667,7 @@ class State:
             del self.modified_last[key]
         except KeyError:
             pass
+
+
+    def count(self):
+        return self.__c
diff --git a/tests/test_state.py b/tests/test_state.py
index 6bbd238..a64c060 100644
--- a/tests/test_state.py
+++ b/tests/test_state.py
@@ -292,6 +292,15 @@ class TestState(unittest.TestCase):
         self.assertIsInstance(v, list)
 
 
+    def test_count(self):
+        states = State(3)
+        states.add('foo')
+        states.add('bar')
+        self.assertEqual(states.count(), 2)
+        states.add('baz')
+        self.assertEqual(states.count(), 3)
+
+
     def test_pure(self):
         states = State(2)
         states.add('foo')
@@ -305,5 +314,13 @@ class TestState(unittest.TestCase):
         self.assertTrue(v)
 
 
+    def test_default(self):
+        states = State(2, default_state='FOO')
+        with self.assertRaises(StateItemNotFound):
+            states.state('NEW')
+        getattr(states, 'FOO')
+        states.state('FOO')
+
+
 if __name__ == '__main__':
     unittest.main()