use generator typehint

delta-io · Nov 12, 2023 · 850a05a · 850a05a
1 parent 2337bd7
commit 850a05a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/deltalake/schema.py b/python/deltalake/schema.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING, Generator, Union
 
 import pyarrow as pa
 import pyarrow.dataset as ds
@@ -82,7 +82,7 @@ def struct_to_delta_dtype(dtype: pa.StructType) -> pa.StructType:
 
 def _cast_schema_to_recordbatchreader(
     reader: pa.RecordBatchReader, schema: pa.schema
-) -> pa.RecordBatch:
+) -> Generator[pa.RecordBatch, None, None]:
     """Creates recordbatch generator."""
     for batch in reader:
         yield pa.Table.from_batches([batch]).cast(schema).to_batches()[0]