Fix various type hints (#1903)

* Fix type hints

* Minor cleanup

Author: rayokota

pyproject.toml

@@ -24,6 +24,9 @@ email = "[email protected]"
 [project.urls]
 Homepage = "https://github.com/confluentinc/confluent-kafka-python"
 
+[tool.mypy]
+ignore_missing_imports = true
+
 [tool.setuptools]
 include-package-data = false
 

src/confluent_kafka/schema_registry/json_schema.py

@@ -78,7 +78,7 @@ def _retrieve_via_httpx(uri: str):
 
 def _resolve_named_schema(
     schema: Schema, schema_registry_client: SchemaRegistryClient,
-    ref_registry: Registry = None
+    ref_registry: Optional[Registry] = None
 ) -> Registry:
     """
     Resolves named schemas referenced by the provided schema recursively.
@@ -222,10 +222,10 @@ def __init__(
         self,
         schema_str: Union[str, Schema, None],
         schema_registry_client: SchemaRegistryClient,
-        to_dict: Callable[[object, SerializationContext], dict] = None,
-        conf: dict = None,
-        rule_conf: dict = None,
-        rule_registry: RuleRegistry = None
+        to_dict: Optional[Callable[[object, SerializationContext], dict]] = None,
+        conf: Optional[dict] = None,
+        rule_conf: Optional[dict] = None,
+        rule_registry: Optional[RuleRegistry] = None
     ):
         super().__init__()
         if isinstance(schema_str, str):
@@ -296,7 +296,7 @@ def __init__(
             rule.configure(self._registry.config() if self._registry else {},
                            rule_conf if rule_conf else {})
 
-    def __call__(self, obj: object, ctx: SerializationContext = None) -> Optional[bytes]:
+    def __call__(self, obj: object, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
         """
         Serializes an object to JSON, prepending it with Confluent Schema Registry
         framing.
@@ -452,11 +452,11 @@ class JSONDeserializer(BaseDeserializer):
     def __init__(
         self,
         schema_str: Union[str, Schema, None],
-        from_dict: Callable[[dict, SerializationContext], object] = None,
-        schema_registry_client: SchemaRegistryClient = None,
-        conf: dict = None,
-        rule_conf: dict = None,
-        rule_registry: RuleRegistry = None
+        from_dict: Optional[Callable[[dict, SerializationContext], object]] = None,
+        schema_registry_client: Optional[SchemaRegistryClient] = None,
+        conf: Optional[dict] = None,
+        rule_conf: Optional[dict] = None,
+        rule_registry: Optional[RuleRegistry] = None
     ):
         super().__init__()
         if isinstance(schema_str, str):
@@ -520,7 +520,7 @@ def __init__(
             rule.configure(self._registry.config() if self._registry else {},
                            rule_conf if rule_conf else {})
 
-    def __call__(self, data: bytes, ctx: SerializationContext = None) -> Union[dict, object, None]:
+    def __call__(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Union[dict, object, None]:
         """
         Deserialize a JSON encoded record with Confluent Schema Registry framing to
         a dict, or object instance according to from_dict if from_dict is specified.
@@ -703,7 +703,7 @@ def _transform_field(
 
 def _validate_subschemas(
     subschemas: List[JsonSchema],
-    message: str
+    message: JsonMessage
 ) -> Optional[JsonSchema]:
     for subschema in subschemas:
         try:

src/confluent_kafka/schema_registry/mock_schema_registry_client.py

@@ -171,8 +171,8 @@ def register_schema_full_response(
         return registered_schema
 
     def get_schema(
-        self, schema_id: int, subject_name: str = None,
-        fmt: str = None
+        self, schema_id: int, subject_name: Optional[str] = None,
+        fmt: Optional[str] = None
     ) -> 'Schema':
         schema = self._store.get_schema(schema_id)
         if schema is not None:
@@ -197,7 +197,7 @@ def get_subjects(self) -> List[str]:
     def delete_subject(self, subject_name: str, permanent: bool = False) -> List[int]:
         return self._store.remove_by_subject(subject_name)
 
-    def get_latest_version(self, subject_name: str, fmt: str = None) -> 'RegisteredSchema':
+    def get_latest_version(self, subject_name: str, fmt: Optional[str] = None) -> 'RegisteredSchema':
         registered_schema = self._store.get_latest_version(subject_name)
         if registered_schema is not None:
             return registered_schema
@@ -206,7 +206,7 @@ def get_latest_version(self, subject_name: str, fmt: str = None) -> 'RegisteredS
 
     def get_latest_with_metadata(
         self, subject_name: str, metadata: Dict[str, str],
-        deleted: bool = False, fmt: str = None
+        deleted: bool = False, fmt: Optional[str] = None
     ) -> 'RegisteredSchema':
         registered_schema = self._store.get_latest_with_metadata(subject_name, metadata)
         if registered_schema is not None:
@@ -216,7 +216,7 @@ def get_latest_with_metadata(
 
     def get_version(
         self, subject_name: str, version: int,
-        deleted: bool = False, fmt: str = None
+        deleted: bool = False, fmt: Optional[str] = None
     ) -> 'RegisteredSchema':
         registered_schema = self._store.get_version(subject_name, version)
         if registered_schema is not None:

src/confluent_kafka/schema_registry/protobuf.py

@@ -187,7 +187,7 @@ def _resolve_named_schema(
     schema: Schema,
     schema_registry_client: SchemaRegistryClient,
     pool: DescriptorPool,
-    visited: Set[str] = None
+    visited: Optional[Set[str]] = None
 ):
     """
     Resolves named schemas referenced by the provided schema recursively.
@@ -373,9 +373,9 @@ def __init__(
         self,
         msg_type: Message,
         schema_registry_client: SchemaRegistryClient,
-        conf: dict = None,
-        rule_conf: dict = None,
-        rule_registry: RuleRegistry = None
+        conf: Optional[dict] = None,
+        rule_conf: Optional[dict] = None,
+        rule_registry: Optional[RuleRegistry] = None
     ):
         super().__init__()
 
@@ -528,7 +528,7 @@ def _resolve_dependencies(
                                                reference.version))
         return schema_refs
 
-    def __call__(self, message: Message, ctx: SerializationContext = None) -> Optional[bytes]:
+    def __call__(self, message: Message, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
         """
         Serializes an instance of a class derived from Protobuf Message, and prepends
         it with Confluent Schema Registry framing.
@@ -675,10 +675,10 @@ class ProtobufDeserializer(BaseDeserializer):
     def __init__(
         self,
         message_type: Message,
-        conf: dict = None,
-        schema_registry_client: SchemaRegistryClient = None,
-        rule_conf: dict = None,
-        rule_registry: RuleRegistry = None
+        conf: Optional[dict] = None,
+        schema_registry_client: Optional[SchemaRegistryClient] = None,
+        rule_conf: Optional[dict] = None,
+        rule_registry: Optional[RuleRegistry] = None
     ):
         super().__init__()
 
@@ -812,7 +812,7 @@ def _read_index_array(buf: io.BytesIO, zigzag: bool = True) -> List[int]:
 
         return msg_index
 
-    def __call__(self, data: bytes, ctx: SerializationContext = None) -> Optional[Message]:
+    def __call__(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Optional[Message]:
         """
         Deserialize a serialized protobuf message with Confluent Schema Registry
         framing.

src/confluent_kafka/schema_registry/rules/encryption/dek_registry/dek_registry_client.py

@@ -471,7 +471,7 @@ def config(self):
 
     def register_kek(
         self, name: str, kms_type: str, kms_key_id: str,
-        shared: bool = False, kms_props: Dict[str, str] = None, doc: str = None
+        shared: bool = False, kms_props: Optional[Dict[str, str]] = None, doc: Optional[str] = None
     ) -> Kek:
         """
         Register a new Key Encryption Key (KEK) with the DEK Registry.

src/confluent_kafka/schema_registry/rules/encryption/dek_registry/mock_dek_registry_client.py

@@ -17,7 +17,7 @@
 #
 
 import time
-from typing import Dict
+from typing import Dict, Optional
 
 from confluent_kafka.schema_registry import SchemaRegistryError
 from confluent_kafka.schema_registry.rules.encryption.dek_registry.dek_registry_client import \
@@ -39,7 +39,7 @@ def __init__(self, conf: dict):
 
     def register_kek(
         self, name: str, kms_type: str, kms_key_id: str,
-        shared: bool = False, kms_props: Dict[str, str] = None, doc: str = None
+        shared: bool = False, kms_props: Optional[Dict[str, str]] = None, doc: Optional[str] = None
     ) -> Kek:
         cache_key = KekId(name=name, deleted=False)
         kek = self._kek_cache.get_kek(cache_key)

src/confluent_kafka/schema_registry/schema_registry_client.py

@@ -175,16 +175,16 @@ def __init__(self, conf: dict):
             raise ValueError("Unrecognized properties: {}"
                              .format(", ".join(conf_copy.keys())))
 
-    def get(self, url: str, query: dict = None) -> Any:
+    def get(self, url: str, query: Optional[dict] = None) -> Any:
         raise NotImplementedError()
 
-    def post(self, url: str, body: dict, **kwargs) -> Any:
+    def post(self, url: str, body: Optional[dict], **kwargs) -> Any:
         raise NotImplementedError()
 
     def delete(self, url: str) -> Any:
         raise NotImplementedError()
 
-    def put(self, url: str, body: dict = None) -> Any:
+    def put(self, url: str, body: Optional[dict] = None) -> Any:
         raise NotImplementedError()
 
 
@@ -209,21 +209,21 @@ def __init__(self, conf: dict):
             timeout=self.timeout
         )
 
-    def get(self, url: str, query: dict = None) -> Any:
+    def get(self, url: str, query: Optional[dict] = None) -> Any:
         return self.send_request(url, method='GET', query=query)
 
-    def post(self, url: str, body: dict, **kwargs) -> Any:
+    def post(self, url: str, body: Optional[dict], **kwargs) -> Any:
         return self.send_request(url, method='POST', body=body)
 
     def delete(self, url: str) -> Any:
         return self.send_request(url, method='DELETE')
 
-    def put(self, url: str, body: dict = None) -> Any:
+    def put(self, url: str, body: Optional[dict] = None) -> Any:
         return self.send_request(url, method='PUT', body=body)
 
     def send_request(
-        self, url: str, method: str, body: dict = None,
-        query: dict = None
+        self, url: str, method: str, body: Optional[dict] = None,
+        query: Optional[dict] = None
     ) -> Any:
         """
         Sends HTTP request to the SchemaRegistry, trying each base URL in turn.
@@ -284,8 +284,8 @@ def send_request(
                                       + str(response.content))
 
     def send_http_request(
-        self, base_url: str, url: str, method: str, headers: dict,
-        body: Optional[str] = None, query: dict = None
+        self, base_url: str, url: str, method: str, headers: Optional[dict],
+        body: Optional[str] = None, query: Optional[dict] = None
     ) -> Response:
         """
         Sends HTTP request to the SchemaRegistry.
@@ -312,6 +312,7 @@ def send_http_request(
         Returns:
             Response: Schema Registry response content.
         """
+        response = None
         for i in range(self.max_retries + 1):
             response = self.session.request(
                 method, url="/".join([base_url, url]),
@@ -324,6 +325,7 @@ def send_http_request(
                 return response
 
             time.sleep(full_jitter(self.retries_wait_ms, self.retries_max_wait_ms, i) / 1000)
+        return response
 
 
 def is_success(status_code: int) -> bool:
@@ -691,7 +693,7 @@ def register_schema_full_response(
         return registered_schema
 
     def get_schema(
-        self, schema_id: int, subject_name: str = None, fmt: str = None
+        self, schema_id: int, subject_name: Optional[str] = None, fmt: Optional[str] = None
     ) -> 'Schema':
         """
         Fetches the schema associated with ``schema_id`` from the
@@ -817,7 +819,7 @@ def delete_subject(self, subject_name: str, permanent: bool = False) -> List[int
         return versions
 
     def get_latest_version(
-        self, subject_name: str, fmt: str = None
+        self, subject_name: str, fmt: Optional[str] = None
     ) -> 'RegisteredSchema':
         """
         Retrieves latest registered version for subject
@@ -853,7 +855,7 @@ def get_latest_version(
 
     def get_latest_with_metadata(
         self, subject_name: str, metadata: Dict[str, str],
-        deleted: bool = False, fmt: str = None
+        deleted: bool = False, fmt: Optional[str] = None
     ) -> 'RegisteredSchema':
         """
         Retrieves latest registered version for subject with the given metadata
@@ -893,7 +895,7 @@ def get_latest_with_metadata(
 
     def get_version(
         self, subject_name: str, version: int,
-        deleted: bool = False, fmt: str = None
+        deleted: bool = False, fmt: Optional[str] = None
     ) -> 'RegisteredSchema':
         """
         Retrieves a specific schema registered under ``subject_name``.
@@ -981,7 +983,7 @@ def delete_version(self, subject_name: str, version: int, permanent: bool = Fals
 
         return response
 
-    def set_compatibility(self, subject_name: Optional[str] = None, level: str = None) -> str:
+    def set_compatibility(self, subject_name: Optional[str] = None, level: Optional[str] = None) -> str:
         """
         Update global or subject level compatibility level.
 
@@ -1068,7 +1070,10 @@ def test_compatibility(
 
         return response['is_compatible']
 
-    def set_config(self, subject_name: Optional[str] = None, config: 'ServerConfig' = None) -> 'ServerConfig':
+    def set_config(
+        self, subject_name: Optional[str] = None,
+        config: Optional['ServerConfig'] = None
+    ) -> 'ServerConfig':
         """
         Update global or subject config.
 

src/confluent_kafka/schema_registry/serde.py

@@ -272,7 +272,7 @@ class BaseSerde(object):
                  '_registry', '_rule_registry', '_subject_name_func',
                  '_field_transformer']
 
-    def _get_reader_schema(self, subject: str, fmt: str = None) -> Optional[RegisteredSchema]:
+    def _get_reader_schema(self, subject: str, fmt: Optional[str] = None) -> Optional[RegisteredSchema]:
         if self._use_latest_with_metadata is not None:
             return self._registry.get_latest_with_metadata(
                 subject, self._use_latest_with_metadata, True, fmt)
@@ -427,6 +427,7 @@ def _has_rules(self, rule_set: RuleSet, mode: RuleMode) -> bool:
                        for rule in rule_set.domain_rules or [])
         elif mode == RuleMode.WRITEREAD:
             return any(rule.mode == mode for rule in rule_set.migration_rules or [])
+        return False
 
     def _get_migrations(
         self, subject: str, source_info: Schema,
@@ -464,7 +465,7 @@ def _get_migrations(
 
     def _get_schemas_between(
         self, subject: str, first: RegisteredSchema,
-        last: RegisteredSchema, fmt: str = None
+        last: RegisteredSchema, fmt: Optional[str] = None
     ) -> List[RegisteredSchema]:
         if last.version - first.version <= 1:
             return [first, last]