From d9d0af994c6edb889e0e264f42ff6aff3444b6b3 Mon Sep 17 00:00:00 2001 From: boyska Date: Mon, 17 Jan 2022 00:21:13 +0100 Subject: [PATCH] typing and cleanup --- larigira/db.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/larigira/db.py b/larigira/db.py index 13aacf0..7aaa40c 100644 --- a/larigira/db.py +++ b/larigira/db.py @@ -3,6 +3,8 @@ from tinydb import TinyDB from tinydb.storages import JSONStorage from tinydb.middlewares import Middleware from pathlib import Path +from typing import Union, Tuple + class ReadOnlyMiddleware(Middleware): """ @@ -20,6 +22,7 @@ class ReadOnlyMiddleware(Middleware): class ReadOnlyException(ValueError): pass + class EventModel(object): def __init__(self, uri, additional_db_dir=None): self.uri = uri @@ -60,14 +63,14 @@ class EventModel(object): self._actions = self._dbs['main'].table("actions") self._alarms = self._dbs['main'].table("alarms") - def canonicalize(self, eid_or_aid): + def canonicalize(self, eid_or_aid: Union[str, int]) -> str: try: int(eid_or_aid) except ValueError: return eid_or_aid return 'main:%d' % eid_or_aid - def parse_id(self, eid_or_aid): + def parse_id(self, eid_or_aid: Union[str, int]) -> Tuple[str, int]: try: int(eid_or_aid) except ValueError: @@ -79,9 +82,9 @@ class EventModel(object): return (dbname, int(num)) - def get_action_by_id(self, action_id): - db, action_id = self.parse_id(action_id) - return self._dbs[db].table('actions').get(eid=action_id) + def get_action_by_id(self, action_id: Union[str, int]): + db, db_action_id = self.parse_id(action_id) + return self._dbs[db].table('actions').get(eid=db_action_id) def get_alarm_by_id(self, alarm_id): db, alarm_id = self.parse_id(alarm_id)