Fix MyPy errors in leveldb

airflow/providers/google/leveldb/hooks/leveldb.py

@@ -23,6 +23,8 @@
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
 
+DB_NOT_INITIALIZED_BEFORE = "The `get_conn` method should be called before!"
+
 
 class LevelDBHookException(AirflowException):
     """Exception specific for LevelDB"""
@@ -43,7 +45,7 @@ def __init__(self, leveldb_conn_id: str = default_conn_name):
         super().__init__()
         self.leveldb_conn_id = leveldb_conn_id
         self.connection = self.get_connection(leveldb_conn_id)
-        self.db = None
+        self.db: Optional[plyvel.DB] = None
 
     def get_conn(self, name: str = '/tmp/testdb/', create_if_missing: bool = False, **kwargs) -> DB:
         """
@@ -74,9 +76,9 @@ def run(
         self,
         command: str,
         key: bytes,
-        value: bytes = None,
-        keys: List[bytes] = None,
-        values: List[bytes] = None,
+        value: Optional[bytes] = None,
+        keys: Optional[List[bytes]] = None,
+        values: Optional[List[bytes]] = None,
     ) -> Optional[bytes]:
         """
         Execute operation with leveldb
@@ -87,21 +89,27 @@ def run(
         :param key: key for command(put,get,delete) execution(, e.g. ``b'key'``, ``b'another-key'``)
         :type key: bytes
         :param value: value for command(put) execution(bytes, e.g. ``b'value'``, ``b'another-value'``)
-        :type value: bytes
+        :type value: Optional[bytes]
         :param keys: keys for command(write_batch) execution(List[bytes], e.g. ``[b'key', b'another-key'])``
-        :type keys: List[bytes]
+        :type keys: Optional[List[bytes]]
         :param values: values for command(write_batch) execution e.g. ``[b'value'``, ``b'another-value']``
-        :type values: List[bytes]
+        :type values: Optional[List[bytes]]
         :returns: value from get or None
         :rtype: Optional[bytes]
         """
         if command == 'put':
+            if not value:
+                raise Exception("Please provide `value`!")
             return self.put(key, value)
         elif command == 'get':
             return self.get(key)
         elif command == 'delete':
             return self.delete(key)
         elif command == 'write_batch':
+            if not keys:
+                raise Exception("Please provide `keys`!")
+            if not values:
+                raise Exception("Please provide `values`!")
             return self.write_batch(keys, values)
         else:
             raise LevelDBHookException("Unknown command for LevelDB hook")
@@ -115,6 +123,8 @@ def put(self, key: bytes, value: bytes):
         :param value: value for put execution e.g. ``b'value'``, ``b'another-value'``
         :type value: bytes
         """
+        if not self.db:
+            raise Exception(DB_NOT_INITIALIZED_BEFORE)
         self.db.put(key, value)
 
     def get(self, key: bytes) -> bytes:
@@ -126,6 +136,8 @@ def get(self, key: bytes) -> bytes:
         :returns: value of key from db.get
         :rtype: bytes
         """
+        if not self.db:
+            raise Exception(DB_NOT_INITIALIZED_BEFORE)
         return self.db.get(key)
 
     def delete(self, key: bytes):
@@ -135,6 +147,8 @@ def delete(self, key: bytes):
         :param key: key for delete execution, e.g. ``b'key'``, ``b'another-key'``
         :type key: bytes
         """
+        if not self.db:
+            raise Exception(DB_NOT_INITIALIZED_BEFORE)
         self.db.delete(key)
 
     def write_batch(self, keys: List[bytes], values: List[bytes]):
@@ -146,6 +160,8 @@ def write_batch(self, keys: List[bytes], values: List[bytes]):
         :param values: values for write_batch execution e.g. ``[b'value', b'another-value']``
         :type values: List[bytes]
         """
+        if not self.db:
+            raise Exception(DB_NOT_INITIALIZED_BEFORE)
         with self.db.write_batch() as batch:
             for i, key in enumerate(keys):
                 batch.put(key, values[i])

airflow/providers/google/leveldb/operators/leveldb.py

@@ -34,11 +34,11 @@ class LevelDBOperator(BaseOperator):
         :param key: key for command(put,get,delete) execution(, e.g. ``b'key'``, ``b'another-key'``)
         :type key: bytes
         :param value: value for command(put) execution(bytes, e.g. ``b'value'``, ``b'another-value'``)
-        :type value: bytes
+        :type value: Optional[bytes]
         :param keys: keys for command(write_batch) execution(List[bytes], e.g. ``[b'key', b'another-key'])``
-        :type keys: List[bytes]
+        :type keys: Optional[List[bytes]]
         :param values: values for command(write_batch) execution e.g. ``[b'value'``, ``b'another-value']``
-        :type values: List[bytes]
+        :type values: Optional[List[bytes]]
         :param leveldb_conn_id:
         :type leveldb_conn_id: str
         :param create_if_missing: whether a new database should be created if needed
@@ -53,9 +53,9 @@ def __init__(
         *,
         command: str,
         key: bytes,
-        value: bytes = None,
-        keys: List[bytes] = None,
-        values: List[bytes] = None,
+        value: Optional[bytes] = None,
+        keys: Optional[List[bytes]] = None,
+        values: Optional[List[bytes]] = None,
         leveldb_conn_id: str = 'leveldb_default',
         name: str = '/tmp/testdb/',
         create_if_missing: bool = True,
@@ -94,5 +94,5 @@ def execute(self, context) -> Optional[str]:
         )
         self.log.info("Done. Returned value was: %s", str(value))
         leveldb_hook.close_conn()
-        value = value if value is None else value.decode()
-        return value
+        str_value = value if value is None else value.decode()
+        return str_value