Fix up some mypy errors (#76)

Co-authored-by: James Bourbeau <jrbourbeau@gmail.com>
Co-authored-by: Miles Granger <miles59923@gmail.com>

.pre-commit-config.yaml

@@ -32,7 +32,7 @@ repos:
           # Type stubs
           - boto3-stubs
           - dask
-          - deltalake
+          - deltalake>=0.16
           - pandas-stubs
           - pytest
           - types-setuptools

dask_deltatable/core.py

@@ -2,7 +2,6 @@
 
 import os
 from collections.abc import Sequence
-from functools import partial
 from typing import Any, Callable, cast
 
 import dask
@@ -121,14 +120,18 @@ def _read_from_filesystem(
     if columns:
         meta = meta[columns]
 
-    kws = dict(meta=meta, label="read-delta-table")
     if not dd._dask_expr_enabled():
         # Setting token not supported in dask-expr
-        kws["token"] = tokenize(path, fs_token, **kwargs)  # type: ignore
+        kwargs["token"] = tokenize(path, fs_token, **kwargs)  # type: ignore
     return dd.from_map(
-        partial(_read_delta_partition, fs=fs, columns=columns, schema=schema, **kwargs),
+        _read_delta_partition,
         pq_files,
-        **kws,
+        fs=fs,
+        columns=columns,
+        schema=schema,
+        meta=meta,
+        label="read-delta-table",
+        **kwargs,
     )
 
 

dask_deltatable/write.py

@@ -17,10 +17,10 @@
 from deltalake import DeltaTable
 
 try:
-    from deltalake.writer import MAX_SUPPORTED_WRITER_VERSION  # type: ignore
+    from deltalake.writer import MAX_SUPPORTED_PYARROW_WRITER_VERSION
 except ImportError:
-    from deltalake.writer import (
-        MAX_SUPPORTED_PYARROW_WRITER_VERSION as MAX_SUPPORTED_WRITER_VERSION,
+    from deltalake.writer import (  # type: ignore
+        MAX_SUPPORTED_WRITER_VERSION as MAX_SUPPORTED_PYARROW_WRITER_VERSION,
     )
 
 from deltalake.writer import (
@@ -177,11 +177,11 @@ def to_deltalake(
         else:
             partition_by = table.metadata().partition_columns
 
-        if table.protocol().min_writer_version > MAX_SUPPORTED_WRITER_VERSION:
+        if table.protocol().min_writer_version > MAX_SUPPORTED_PYARROW_WRITER_VERSION:
             raise DeltaProtocolError(
                 "This table's min_writer_version is "
                 f"{table.protocol().min_writer_version}, "
-                f"but this method only supports version {MAX_SUPPORTED_WRITER_VERSION}."
+                f"but this method only supports version {MAX_SUPPORTED_PYARROW_WRITER_VERSION}."
             )
     else:  # creating a new table
         current_version = -1

requirements.txt

@@ -1,4 +1,4 @@
 dask[dataframe]
-deltalake>=0.15
+deltalake>=0.16
 fsspec
 pyarrow

tests/test_acceptance.py

@@ -60,9 +60,8 @@ def test_reader_all_primitive_types():
     # Dask and delta go through different parquet parsers which read the
     # timestamp differently. This is likely a bug in arrow but the delta result
     # is "more correct".
-    expected_ddf["timestamp"] = (
-        expected_ddf["timestamp"].astype("datetime64[us]").dt.tz_localize("UTC")
-    )
+    expected_ddf["timestamp"] = expected_ddf["timestamp"].astype("datetime64[us]")
+    expected_ddf["timestamp"] = expected_ddf["timestamp"].dt.tz_localize("UTC")
     assert_eq(actual_ddf, expected_ddf)