GSK-2874 Fix quality badge using https://peps.python.org/pep-0563/

giskard/core/suite.py

@@ -1,4 +1,6 @@
-from typing import Any, Callable, Dict, List, Optional, Union
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 import inspect
 import logging
@@ -28,6 +30,9 @@
 from ..utils.artifacts import serialize_parameter
 from .kwargs_utils import get_imports_code
 
+if TYPE_CHECKING:
+    from mlflow import MlflowClient
+
 logger = logging.getLogger(__name__)
 
 suite_input_types: List[type] = [
@@ -123,10 +128,9 @@ def _repr_html_(self):
         widget = TestSuiteResultWidget(self)
         return widget.render_html()
 
-    def to_mlflow(self, mlflow_client=None, mlflow_run_id: str = None):
+    def to_mlflow(self, mlflow_client: MlflowClient = None, mlflow_run_id: str = None):
         import mlflow
 
-        mlflow_client: mlflow.MlflowClient = mlflow_client
         from giskard.integrations.mlflow.giskard_evaluator_utils import process_text
 
         metrics = dict()

giskard/datasets/base/__init__.py

@@ -1,4 +1,6 @@
-from typing import Dict, Hashable, List, Optional, Union
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Hashable, List, Optional, Union
 
 import inspect
 import logging
@@ -33,6 +35,9 @@
 from ...utils.file_utils import get_file_name
 from ..metadata.indexing import ColumnMetadataMixin
 
+if TYPE_CHECKING:
+    from mlflow import MlflowClient
+
 SAMPLE_SIZE = 1000
 
 logger = logging.getLogger(__name__)
@@ -692,10 +697,9 @@ def copy(self):
 
         return dataset
 
-    def to_mlflow(self, mlflow_client=None, mlflow_run_id: str = None):
+    def to_mlflow(self, mlflow_client: MlflowClient = None, mlflow_run_id: str = None):
         import mlflow
 
-        mlflow_client: mlflow.MlflowClient = mlflow_client  # Doing typing here, to avoid import from mlflow
         # To avoid file being open in write mode and read at the same time,
         # First, we'll write it, then make sure to remove it
         with tempfile.NamedTemporaryFile(prefix="dataset-", suffix=".csv", delete=False) as f:

giskard/scanner/report.py

@@ -1,4 +1,6 @@
-from typing import Optional
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
 
 import random
 import string
@@ -11,6 +13,9 @@
 from giskard.core.errors import GiskardImportError
 from giskard.utils.analytics_collector import analytics, anonymize
 
+if TYPE_CHECKING:
+    from mlflow import MlflowClient
+
 
 class ScanReport:
     def __init__(self, issues, model=None, dataset=None, as_html: bool = True):
@@ -212,7 +217,7 @@ def get_scan_summary_for_mlflow(scan_results):
 
     def to_mlflow(
         self,
-        mlflow_client=None,
+        mlflow_client: MlflowClient = None,
         mlflow_run_id: str = None,
         summary: bool = True,
         model_artifact_path: str = "",
@@ -223,7 +228,6 @@ def to_mlflow(
         """
         import mlflow
 
-        mlflow_client: mlflow.MlflowClient = mlflow_client
         results_df = self.get_scan_summary_for_mlflow(self)
         if model_artifact_path != "":
             model_artifact_path = "-for-" + model_artifact_path