Run MyPy on CDK/base-python and fix issues. (#3175)

airbyte-cdk/python/airbyte_cdk/base_python/integration.py

@@ -25,7 +25,7 @@ import json
 import os
 import pkgutil
 from collections import defaultdict
-from typing import Dict, Generator
+from typing import Dict, Iterable, Any, Optional, Mapping, MutableMapping
 
 from airbyte_cdk.models import AirbyteCatalog, AirbyteConnectionStatus, AirbyteMessage, ConfiguredAirbyteCatalog, ConnectorSpecification
 
@@ -34,8 +34,8 @@ from .logger import AirbyteLogger
 
 class AirbyteSpec(object):
     @staticmethod
-    def from_file(file: str):
-        with open(file) as file:
+    def from_file(file_name: str):
+        with open(file_name) as file:
             spec_text = file.read()
         return AirbyteSpec(spec_text)
 
@@ -46,7 +46,7 @@ class AirbyteSpec(object):
 class Integration(object):
 
     # can be overridden to change an input config
-    def configure(self, config: json, temp_dir: str) -> json:
+    def configure(self, config: Mapping[str, Any], temp_dir: str) -> Mapping[str, Any]:
         """
         Persist config in temporary directory to run the Source job
         """
@@ -55,13 +55,13 @@ class Integration(object):
         return config
 
     @staticmethod
-    def read_config(config_path: str) -> json:
+    def read_config(config_path: str) -> Mapping[str, Any]:
         with open(config_path, "r") as file:
             contents = file.read()
         return json.loads(contents)
 
     @staticmethod
-    def write_config(config: json, config_path: str):
+    def write_config(config: Mapping[str, Any], config_path: str):
         with open(config_path, "w") as fh:
             fh.write(json.dumps(config))
 
@@ -70,7 +70,7 @@ class Integration(object):
         return ConfiguredAirbyteCatalog.parse_obj(self.read_config(catalog_path))
 
     # can be overridden to change an input state
-    def read_state(self, state_path: str) -> Dict[str, any]:
+    def read_state(self, state_path: str) -> Dict[str, Any]:
         if state_path:
             state_obj = json.loads(open(state_path, "r").read())
         else:
@@ -83,17 +83,19 @@ class Integration(object):
         Returns the spec for this integration. The spec is a JSON-Schema object describing the required configurations (e.g: username and password)
         required to run this integration.
         """
-        raw_spec = pkgutil.get_data(self.__class__.__module__.split(".")[0], "spec.json")
+        raw_spec: Optional[bytes] = pkgutil.get_data(self.__class__.__module__.split(".")[0], "spec.json")
+        if not raw_spec:
+            raise ValueError("Unable to find spec.json.")
         return ConnectorSpecification.parse_obj(json.loads(raw_spec))
 
-    def check(self, logger: AirbyteLogger, config: json) -> AirbyteConnectionStatus:
+    def check(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> AirbyteConnectionStatus:
         """
         Tests if the input configuration can be used to successfully connect to the integration e.g: if a provided Stripe API token can be used to connect
         to the Stripe API.
         """
         raise Exception("Not Implemented")
 
-    def discover(self, logger: AirbyteLogger, config: json) -> AirbyteCatalog:
+    def discover(self, logger: AirbyteLogger, config: Mapping[str, Any]) -> AirbyteCatalog:
         """
         Returns an AirbyteCatalog representing the available streams and fields in this integration. For example, given valid credentials to a
         Postgres database, returns an Airbyte catalog where each postgres table is a stream, and each table column is a field.
@@ -106,8 +108,8 @@ class Source(Integration):
         super().__init__()
 
     def read(
-        self, logger: AirbyteLogger, config: json, catalog: ConfiguredAirbyteCatalog, state_path: Dict[str, any]
-    ) -> Generator[AirbyteMessage, None, None]:
+        self, logger: AirbyteLogger, config: Mapping[str, Any], catalog: ConfiguredAirbyteCatalog, state: MutableMapping[str, Any] = None
+    ) -> Iterable[AirbyteMessage]:
         """
         Returns a generator of the AirbyteMessages generated by reading the source with the given configuration, catalog, and state.
         """