add types for SR rules

Author: fangnx

src/confluent_kafka/schema_registry/error.py

@@ -15,6 +15,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+from typing import Optional
+
 try:
     from fastavro.schema import SchemaParseException, UnknownType
 except ImportError:
@@ -41,15 +43,15 @@ class SchemaRegistryError(Exception):
     """  # noqa: E501
     UNKNOWN = -1
 
-    def __init__(self, http_status_code: int, error_code: int, error_message: str):
+    def __init__(self, http_status_code: int, error_code: int, error_message: str) -> None:
         self.http_status_code = http_status_code
         self.error_code = error_code
         self.error_message = error_message
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return str(self)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return "{} (HTTP status code {}, SR code {})".format(self.error_message,
                                                              self.http_status_code,
                                                              self.error_code)
@@ -58,7 +60,8 @@ def __str__(self):
 class OAuthTokenError(Exception):
     """Raised when an OAuth token cannot be retrieved."""
 
-    def __init__(self, message, status_code=None, response_text=None):
+    def __init__(self, message: str, status_code: Optional[int] = None,
+                 response_text: Optional[str] = None) -> None:
         self.message = message
         self.status_code = status_code
         self.response_text = response_text

src/confluent_kafka/schema_registry/rules/cel/cel_field_presence.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import threading
+from typing import Any
 
 import celpy  # type: ignore
 
@@ -33,9 +34,9 @@ def in_has() -> bool:
 
 
 class InterpretedRunner(celpy.InterpretedRunner):
-    def evaluate(self, context):
+    def evaluate(self, context: Any) -> Any:
         class Evaluator(celpy.Evaluator):
-            def macro_has_eval(self, exprlist) -> celpy.celtypes.BoolType:
+            def macro_has_eval(self, exprlist: Any) -> celpy.celtypes.BoolType:
                 _has_state.in_has = True
                 result = super().macro_has_eval(exprlist)
                 _has_state.in_has = False

src/confluent_kafka/schema_registry/rules/encryption/awskms/aws_driver.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from typing import Dict, Any, Optional
 
 import boto3
 import tink
@@ -34,13 +35,13 @@
 
 
 class AwsKmsDriver(KmsDriver):
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
     def get_key_url_prefix(self) -> str:
         return _PREFIX
 
-    def new_kms_client(self, conf: dict, key_url: str) -> KmsClient:
+    def new_kms_client(self, conf: Dict[str, Any], key_url: Optional[str]) -> KmsClient:
         uri_prefix = _PREFIX
         if key_url is not None:
             uri_prefix = key_url
@@ -92,7 +93,7 @@ def new_kms_client(self, conf: dict, key_url: str) -> KmsClient:
         return new_client(boto3_client=client, key_uri=uri_prefix)
 
     @classmethod
-    def register(cls):
+    def register(cls) -> None:
         register_kms_driver(AwsKmsDriver())
 
 

src/confluent_kafka/schema_registry/rules/encryption/azurekms/azure_driver.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Dict, Any, Optional
 
 from azure.identity import DefaultAzureCredential, ClientSecretCredential
 
@@ -28,13 +29,13 @@
 
 
 class AzureKmsDriver(KmsDriver):
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
     def get_key_url_prefix(self) -> str:
         return _PREFIX
 
-    def new_kms_client(self, conf: dict, key_url: str) -> KmsClient:
+    def new_kms_client(self, conf: Dict[str, Any], key_url: Optional[str]) -> KmsClient:
         uri_prefix = _PREFIX
         if key_url is not None:
             uri_prefix = key_url
@@ -50,5 +51,5 @@ def new_kms_client(self, conf: dict, key_url: str) -> KmsClient:
         return AzureKmsClient(uri_prefix, creds)
 
     @classmethod
-    def register(cls):
+    def register(cls) -> None:
         register_kms_driver(AzureKmsDriver())

src/confluent_kafka/schema_registry/rules/encryption/gcpkms/gcp_driver.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from typing import Dict, Any, Optional
 
 import tink
 from google.oauth2 import service_account
@@ -32,13 +33,13 @@
 
 
 class GcpKmsDriver(KmsDriver):
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
     def get_key_url_prefix(self) -> str:
         return _PREFIX
 
-    def new_kms_client(self, conf: dict, key_url: str) -> KmsClient:
+    def new_kms_client(self, conf: Dict[str, Any], key_url: Optional[str]) -> KmsClient:
         uri_prefix = _PREFIX
         if key_url is not None:
             uri_prefix = key_url
@@ -70,7 +71,7 @@ def new_kms_client(self, conf: dict, key_url: str) -> KmsClient:
         return _GcpKmsClient(uri_prefix, creds)
 
     @classmethod
-    def register(cls):
+    def register(cls) -> None:
         register_kms_driver(GcpKmsDriver())
 
 

src/confluent_kafka/schema_registry/rules/encryption/hcvault/hcvault_driver.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from typing import Dict, Any, Optional
 
 from tink import KmsClient
 
@@ -28,13 +29,13 @@
 
 
 class HcVaultKmsDriver(KmsDriver):
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
     def get_key_url_prefix(self) -> str:
         return _PREFIX
 
-    def new_kms_client(self, conf: dict, key_url: str) -> KmsClient:
+    def new_kms_client(self, conf: Dict[str, Any], key_url: Optional[str]) -> KmsClient:
         uri_prefix = _PREFIX
         if key_url is not None:
             uri_prefix = key_url
@@ -53,5 +54,5 @@ def new_kms_client(self, conf: dict, key_url: str) -> KmsClient:
         return HcVaultKmsClient(uri_prefix, token, namespace, role_id, secret_id)
 
     @classmethod
-    def register(cls):
+    def register(cls) -> None:
         register_kms_driver(HcVaultKmsDriver())

src/confluent_kafka/schema_registry/rules/encryption/localkms/local_driver.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from typing import Dict, Any, Optional
 
 import tink
 from tink import KmsClient
@@ -26,13 +27,13 @@
 
 
 class LocalKmsDriver(KmsDriver):
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
     def get_key_url_prefix(self) -> str:
         return _PREFIX
 
-    def new_kms_client(self, conf: dict, key_url: str) -> KmsClient:
+    def new_kms_client(self, conf: Dict[str, Any], key_url: Optional[str]) -> KmsClient:
         secret = conf.get(_SECRET)
         if secret is None:
             secret = os.getenv("LOCAL_SECRET")
@@ -41,5 +42,5 @@ def new_kms_client(self, conf: dict, key_url: str) -> KmsClient:
         return LocalKmsClient(secret)
 
     @classmethod
-    def register(cls):
+    def register(cls) -> None:
         register_kms_driver(LocalKmsDriver())