1 files changed, 12 insertions, 10 deletions

diff --git a/src/db.py b/src/db.py
index 6f25ec3..5173dda 100755
--- a/src/db.py
+++ b/src/db.py
@@ -7,6 +7,7 @@
 # value if you're too quick with generating the timestamps, ie
 # invoking time.time() several times quickly enough.
 
+from typing import Dict, List, Tuple, Union, Any
 import os
 import sys
 import time
@@ -16,7 +17,7 @@ from schema import as_index_list, validate_collector_data
 
 
 class DictDB():
-    def __init__(self):
+    def __init__(self) -> None:
         """
         Check if the database exists, otherwise we will create it together
         with the indexes specified in index.py.
@@ -35,7 +36,7 @@ class DictDB():
         if 'COUCHDB_PORT' in os.environ:
             couchdb_port = os.environ['COUCHDB_PORT']
         else:
-            couchdb_port = 5984
+            couchdb_port = "5984"
 
         self.server = couch.client.Server(
             f"http://{self.username}:{self.password}@{self.hostname}:{couchdb_port}/")
@@ -52,7 +53,7 @@ class DictDB():
 
         self._ts = time.time()
 
-    def unique_key(self):
+    def unique_key(self) -> int:
         """
         Create a unique key based on the current time. We will use this as
         the ID for any new documents we store in CouchDB.
@@ -65,18 +66,19 @@ class DictDB():
 
         return self._ts
 
-    def add(self, data, batch_write=False):
+    # Why batch_write???
+    def add(self, data: Union[List[Dict[str, Any]], Dict[str, Any]]) -> Union[str, Tuple[str, str]]:
         """
         Store a document in CouchDB.
         """
 
-        if type(data) is list:
+        if isinstance(data, List):
             for item in data:
                 error = validate_collector_data(item)
                 if error != "":
                     return error
                 item['_id'] = str(self.unique_key())
-            ret = self.couchdb.save_bulk(data)
+            ret: Tuple[str, str] = self.couchdb.save_bulk(data)
         else:
             error = validate_collector_data(data)
             if error != "":
@@ -86,13 +88,13 @@ class DictDB():
 
         return ret
 
-    def get(self, key):
+    def get(self, key: int) -> Dict[str, Any]:
         """
         Get a document based on its ID, return an empty dict if not found.
         """
 
         try:
-            doc = self.couchdb.get(key)
+            doc: Dict[str, Any] = self.couchdb.get(key)
         except couch.exceptions.NotFound:
             doc = {}
 
@@ -101,7 +103,7 @@ class DictDB():
     def slice(self, key_from=None, key_to=None):
         pass
 
-    def search(self, limit=25, skip=0, **kwargs):
+    def search(self, limit: int = 25, skip: int = 0, **kwargs: Any) -> List[Dict[str, Any]]:
         """
         Execute a Mango query, ideally we should have an index matching
         the query otherwise things will be slow.
@@ -134,7 +136,7 @@ class DictDB():
 
         return data
 
-    def delete(self, key):
+    def delete(self, key: int) -> Union[int, None]:
         """
         Delete a document based on its ID.
         """