lint

Signed-off-by: Sajid Alam <[email protected]>

Author: SajidAlamQB

Makefile

@@ -46,6 +46,7 @@ sign-off:
 test-no-spark: dataset-doctests-no-spark
 	cd kedro-datasets && pytest tests --no-cov --ignore tests/spark --ignore tests/databricks --numprocesses 4 --dist loadfile
 
+
 # kedro-datasets/snowflake tests skipped from default scope
 test-snowflake-only:
 	cd kedro-datasets && pytest --no-cov --numprocesses 1 --dist loadfile -m snowflake

kedro-datasets/kedro_datasets/_utils/databricks_utils.py

@@ -1,3 +1,7 @@
+"""Utility functions for Databricks."""
+from __future__ import annotations
+
+import logging
 import os
 from fnmatch import fnmatch
 from pathlib import PurePosixPath
@@ -10,8 +14,6 @@
     from databricks.connect import DatabricksSession
     from pyspark.dbutils import DBUtils
 
-import logging
-
 logger = logging.getLogger(__name__)
 
 
@@ -36,7 +38,7 @@ def strip_dbfs_prefix(path: str, prefix: str = "/dbfs") -> str:
     return path[len(prefix) :] if path.startswith(prefix) else path
 
 
-def dbfs_glob(pattern: str, dbutils: "DBUtils") -> list[str]:
+def dbfs_glob(pattern: str, dbutils: DBUtils) -> list[str]:
     """Perform a custom glob search in DBFS using the provided pattern.
     It is assumed that version paths are managed by Kedro only.
 
@@ -63,7 +65,7 @@ def dbfs_glob(pattern: str, dbutils: "DBUtils") -> list[str]:
     return sorted(matched)
 
 
-def get_dbutils(spark: Union[SparkSession, "DatabricksSession"]) -> "DBUtils":
+def get_dbutils(spark: SparkSession | DatabricksSession) -> DBUtils:
     """Get the instance of 'dbutils' or None if the one could not be found."""
     dbutils = globals().get("dbutils")
     if dbutils:
@@ -87,7 +89,7 @@ def get_dbutils(spark: Union[SparkSession, "DatabricksSession"]) -> "DBUtils":
     return dbutils
 
 
-def dbfs_exists(pattern: str, dbutils: "DBUtils") -> bool:
+def dbfs_exists(pattern: str, dbutils: DBUtils) -> bool:
     """Perform an `ls` list operation in DBFS using the provided pattern.
     It is assumed that version paths are managed by Kedro.
     Broad `Exception` is present due to `dbutils.fs.ExecutionError` that

kedro-datasets/kedro_datasets/_utils/spark_utils.py

@@ -1,4 +1,8 @@
+"""Utility functions for Spark."""
+from __future__ import annotations
+
 import json
+import logging
 import os
 from copy import deepcopy
 from pathlib import PurePosixPath
@@ -18,12 +22,10 @@
     from databricks.connect import DatabricksSession
     from pyspark.sql.types import StructType
 
-import logging
-
 logger = logging.getLogger(__name__)
 
 
-def get_spark() -> Union[SparkSession, "DatabricksSession"]:
+def get_spark() -> SparkSession | DatabricksSession:
     """
     Returns the SparkSession. In case databricks-connect is available we use it for
     extended configuration mechanisms and notebook compatibility,
@@ -46,7 +48,7 @@ def get_spark() -> Union[SparkSession, "DatabricksSession"]:
     return spark
 
 
-def get_spark_with_remote_support() -> Union[SparkSession, "DatabricksSession"]:
+def get_spark_with_remote_support() -> SparkSession | DatabricksSession:
     """Get Spark session with support for Spark Connect and Databricks Connect.
 
     This function attempts to create a Spark session in the following order:
@@ -183,7 +185,7 @@ def get_spark_filesystem(
         raise ImportError(msg) from exc
 
 
-def load_spark_schema_from_file(schema_config: dict[str, Any]) -> "StructType":
+def load_spark_schema_from_file(schema_config: dict[str, Any]) -> StructType:
     """Load Spark schema from JSON file.
 
     Args:

kedro-datasets/tests/_utils/test_databricks_util.py

@@ -32,7 +32,7 @@ def test_get_dbutils_no_ipython(self, mock_globals):
 
         with patch.dict("sys.modules", {"pyspark.dbutils": None, "IPython": None}):
             result = get_dbutils(MagicMock())
-            assert result is None or result is not None
+            assert result is None
 
 
 class TestValidateDatabricksPath: