diff --git a/qurator/eynollah/utils/counter.py b/qurator/eynollah/utils/counter.py index 6bd9d7d..bc1d765 100644 --- a/qurator/eynollah/utils/counter.py +++ b/qurator/eynollah/utils/counter.py @@ -7,8 +7,13 @@ class EynollahIdCounter(): def __init__(self, region_idx=0, line_idx=0): self._counter = Counter() - self.set('region', region_idx) - self.set('line', line_idx) + self._inital_region_idx = region_idx + self._inital_line_idx = line_idx + self.reset() + + def reset(self): + self.set('region', self._inital_region_idx) + self.set('line', self._inital_line_idx) def inc(self, name, val=1): self._counter.update({name: val}) diff --git a/tests/test_counter.py b/tests/test_counter.py index a157cd5..8ef0756 100644 --- a/tests/test_counter.py +++ b/tests/test_counter.py @@ -14,6 +14,10 @@ def test_counter_string(): def test_counter_init(): c = EynollahIdCounter(region_idx=2) assert c.get('region') == 2 + c.inc('region') + assert c.get('region') == 3 + c.reset() + assert c.get('region') == 2 def test_counter_methods(): c = EynollahIdCounter()