Fix lint issues

Author: markdlv-sf

src/datacustomcode/io/reader/query_api.py

@@ -21,6 +21,7 @@
     Union,
 )
 
+import pandas.api.types as pd_types
 from pyspark.sql.types import (
     BooleanType,
     DoubleType,
@@ -34,7 +35,6 @@
 
 from datacustomcode.credentials import Credentials
 from datacustomcode.io.reader.base import BaseDataCloudReader
-import pandas.api.types as pd_types
 
 if TYPE_CHECKING:
     import pandas
@@ -58,6 +58,7 @@ def _pandas_to_spark_schema(
 ) -> StructType:
     fields = []
     for column, dtype in pandas_df.dtypes.items():
+        spark_type: AtomicType
         if pd_types.is_datetime64_any_dtype(dtype):
             spark_type = TimestampType()
         else:

tests/io/reader/test_query_api.py

@@ -62,14 +62,23 @@ def test_pandas_to_spark_schema_nullable(self):
 
     def test_pandas_to_spark_schema_datetime_types(self):
         """Test conversion of pandas datetime types to Spark TimestampType."""
-        import numpy as np
 
         # Create test data with different datetime types
         data = {
-            "datetime_ns": pd.to_datetime(["2023-01-01 10:00:00", "2023-01-02 11:00:00"]),
-            "datetime_ns_utc": pd.to_datetime(["2023-01-01 10:00:00", "2023-01-02 11:00:00"], utc=True),
-            "datetime_ms": pd.to_datetime(["2023-01-01 10:00:00", "2023-01-02 11:00:00"]).astype("datetime64[ms]"),
-            "datetime_ms_utc": pd.to_datetime(["2023-01-01 10:00:00", "2023-01-02 11:00:00"], utc=True).tz_localize(None).astype("datetime64[ms]"),
+            "datetime_ns": pd.to_datetime(
+                ["2023-01-01 10:00:00", "2023-01-02 11:00:00"]
+            ),
+            "datetime_ns_utc": pd.to_datetime(
+                ["2023-01-01 10:00:00", "2023-01-02 11:00:00"], utc=True
+            ),
+            "datetime_ms": pd.to_datetime(
+                ["2023-01-01 10:00:00", "2023-01-02 11:00:00"]
+            ).astype("datetime64[ms]"),
+            "datetime_ms_utc": pd.to_datetime(
+                ["2023-01-01 10:00:00", "2023-01-02 11:00:00"], utc=True
+            )
+            .tz_localize(None)
+            .astype("datetime64[ms]"),
         }
         df = pd.DataFrame(data)
 
@@ -82,9 +91,16 @@ def test_pandas_to_spark_schema_datetime_types(self):
 
         # Check that all datetime columns map to TimestampType
         field_dict = {field.name: field for field in schema.fields}
-        for field_name in ["datetime_ns", "datetime_ns_utc", "datetime_ms", "datetime_ms_utc"]:
-            assert isinstance(field_dict[field_name].dataType, TimestampType), \
-                f"Field {field_name} should be TimestampType, got {type(field_dict[field_name].dataType)}"
+        for field_name in [
+            "datetime_ns",
+            "datetime_ns_utc",
+            "datetime_ms",
+            "datetime_ms_utc",
+        ]:
+            assert isinstance(field_dict[field_name].dataType, TimestampType), (
+                f"Field {field_name} should be TimestampType, "
+                f"got {type(field_dict[field_name].dataType)}"
+            )
             assert field_dict[field_name].nullable
 
         # Verify the actual pandas dtypes to ensure our test data has the expected types