fix(ingest/build): Fix sagemaker mypy and flake8 issues (#8530)

Author: treff7es

metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/feature_groups.py

@@ -24,7 +24,7 @@
     from mypy_boto3_sagemaker import SageMakerClient
     from mypy_boto3_sagemaker.type_defs import (
         DescribeFeatureGroupResponseTypeDef,
-        FeatureDefinitionOutputTypeDef,
+        FeatureDefinitionTypeDef,
         FeatureGroupSummaryTypeDef,
     )
 
@@ -147,7 +147,7 @@ def get_feature_type(self, aws_type: str, feature_name: str) -> str:
     def get_feature_wu(
         self,
         feature_group_details: "DescribeFeatureGroupResponseTypeDef",
-        feature: "FeatureDefinitionOutputTypeDef",
+        feature: "FeatureDefinitionTypeDef",
     ) -> MetadataWorkUnit:
         """
         Generate an MLFeature workunit for a SageMaker feature.

metadata-ingestion/src/datahub/utilities/mapping.py

@@ -252,11 +252,11 @@ def sanitize_owner_ids(self, owner_id: str) -> str:
 
     def get_match(self, match_clause: Any, raw_props_value: Any) -> Optional[Match]:
         # function to check if a match clause is satisfied to a value.
-        if type(raw_props_value) not in Constants.OPERAND_DATATYPE_SUPPORTED or type(
-            raw_props_value
-        ) != type(match_clause):
+        if not any(
+            isinstance(raw_props_value, t) for t in Constants.OPERAND_DATATYPE_SUPPORTED
+        ) or not isinstance(raw_props_value, type(match_clause)):
             return None
-        elif type(raw_props_value) == str:
+        elif isinstance(raw_props_value, str):
             return re.match(match_clause, raw_props_value)
         else:
             return re.match(str(match_clause), str(raw_props_value))

metadata-ingestion/src/datahub_provider/_plugin.py

@@ -107,7 +107,7 @@ def get_inlets_from_task(task: BaseOperator, context: Any) -> Iterable[Any]:
         ]
 
         for inlet in task_inlets:
-            if type(inlet) != str:
+            if isinstance(inlet, str):
                 inlets.append(inlet)
 
     return inlets

metadata-ingestion/src/datahub_provider/operators/datahub_assertion_operator.py

@@ -62,9 +62,9 @@ def execute(self, context: Any) -> bool:
             return True
 
         self.log.info(f"Checking if dataset {self.urn} is ready to be consumed")
-        if type(self.urn) == str:
+        if isinstance(self.urn, str):
             urns = [self.urn]
-        elif type(self.urn) == list:
+        elif isinstance(self.urn, list):
             urns = self.urn
         else:
             raise Exception(f"urn parameter has invalid type {type(self.urn)}")

metadata-ingestion/src/datahub_provider/operators/datahub_assertion_sensor.py

@@ -61,9 +61,9 @@ def poke(self, context: Any) -> bool:
             return True
 
         self.log.info(f"Checking if dataset {self.urn} is ready to be consumed")
-        if type(self.urn) == str:
+        if isinstance(self.urn, str):
             urns = [self.urn]
-        elif type(self.urn) == list:
+        elif isinstance(self.urn, list):
             urns = self.urn
         else:
             raise Exception(f"urn parameter has invalid type {type(self.urn)}")

metadata-ingestion/src/datahub_provider/operators/datahub_operation_operator.py

@@ -76,9 +76,9 @@ def execute(self, context: Any) -> bool:
             return True
 
         self.log.info(f"Checking if dataset {self.urn} is ready to be consumed")
-        if type(self.urn) == str:
+        if isinstance(self.urn, str):
             urns = [self.urn]
-        elif type(self.urn) == list:
+        elif isinstance(self.urn, list):
             urns = self.urn
         else:
             raise Exception(f"urn parameter has invalid type {type(self.urn)}")

metadata-ingestion/src/datahub_provider/operators/datahub_operation_sensor.py

@@ -78,9 +78,9 @@ def poke(self, context: Any) -> bool:
             return True
 
         self.log.info(f"Checking if dataset {self.urn} is ready to be consumed")
-        if type(self.urn) == str:
+        if isinstance(self.urn, str):
             urns = [self.urn]
-        elif type(self.urn) == list:
+        elif isinstance(self.urn, list):
             urns = self.urn
         else:
             raise Exception(f"urn parameter has invalid type {type(self.urn)}")

metadata-ingestion/tests/unit/test_glue_source.py

@@ -89,7 +89,7 @@ def test_column_type(hive_column_type: str, expected_type: Type) -> None:
     )
     schema_fields = avro_schema_to_mce_fields(json.dumps(avro_schema))
     actual_schema_field_type = schema_fields[0].type
-    assert type(actual_schema_field_type.type) == expected_type
+    assert isinstance(actual_schema_field_type.type, expected_type)
 
 
 @pytest.mark.parametrize(

metadata-ingestion/tests/unit/test_kafka_source.py

@@ -132,7 +132,9 @@ def test_kafka_source_workunits_with_platform_instance(mock_kafka, mock_admin_cl
 
     # DataPlatform aspect should be present when platform_instance is configured
     data_platform_aspects = [
-        asp for asp in proposed_snap.aspects if type(asp) == DataPlatformInstanceClass
+        asp
+        for asp in proposed_snap.aspects
+        if isinstance(asp, DataPlatformInstanceClass)
     ]
     assert len(data_platform_aspects) == 1
     assert data_platform_aspects[0].instance == make_dataplatform_instance_urn(
@@ -141,7 +143,7 @@ def test_kafka_source_workunits_with_platform_instance(mock_kafka, mock_admin_cl
 
     # The default browse path should include the platform_instance value
     browse_path_aspects = [
-        asp for asp in proposed_snap.aspects if type(asp) == BrowsePathsClass
+        asp for asp in proposed_snap.aspects if isinstance(asp, BrowsePathsClass)
     ]
     assert len(browse_path_aspects) == 1
     assert f"/prod/{PLATFORM}/{PLATFORM_INSTANCE}" in browse_path_aspects[0].paths
@@ -177,13 +179,15 @@ def test_kafka_source_workunits_no_platform_instance(mock_kafka, mock_admin_clie
 
     # DataPlatform aspect should not be present when platform_instance is not configured
     data_platform_aspects = [
-        asp for asp in proposed_snap.aspects if type(asp) == DataPlatformInstanceClass
+        asp
+        for asp in proposed_snap.aspects
+        if isinstance(asp, DataPlatformInstanceClass)
     ]
     assert len(data_platform_aspects) == 0
 
     # The default browse path should include the platform_instance value
     browse_path_aspects = [
-        asp for asp in proposed_snap.aspects if type(asp) == BrowsePathsClass
+        asp for asp in proposed_snap.aspects if isinstance(asp, BrowsePathsClass)
     ]
     assert len(browse_path_aspects) == 1
     assert f"/prod/{PLATFORM}" in browse_path_aspects[0].paths

metadata-ingestion/tests/unit/utilities/test_file_backed_collections.py

@@ -256,7 +256,7 @@ def test_shared_connection() -> None:
         iterator = cache2.sql_query_iterator(
             f"SELECT y, sum(x) FROM {cache2.tablename} GROUP BY y ORDER BY y"
         )
-        assert type(iterator) == sqlite3.Cursor
+        assert isinstance(iterator, sqlite3.Cursor)
         assert [tuple(r) for r in iterator] == [("a", 15), ("b", 11)]
 
         # Test joining between the two tables.