Add missing partitioning typing case (#132)

This should now support the examples in the docstring for partitioning.
zen-xu · Oct 24, 2024 · 090ea66 · 090ea66
1 parent 717e64b
commit 090ea66
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/pyarrow-stubs/_dataset.pyi b/pyarrow-stubs/_dataset.pyi
@@ -351,6 +351,12 @@ class DirectoryPartitioning(KeyValuePartitioning):
         schema: lib.Schema | None = None,
         segment_encoding: Literal["uri", "none"] = "uri",
     ) -> PartitioningFactory: ...
+    def __init__(
+        self,
+        schema: lib.Schema,
+        dictionaries: dict[str, lib.Array] | None = None,
+        segment_encoding: Literal["uri", "none"] = "uri",
+    ) -> None: ...
 
 class HivePartitioning(KeyValuePartitioning):
     def __init__(

diff --git a/pyarrow-stubs/dataset.pyi b/pyarrow-stubs/dataset.pyi
@@ -111,6 +111,10 @@ __all__ = [
 
 _DatasetFormat: TypeAlias = Literal["parquet", "ipc", "arrow", "feather", "csv"]
 
+@overload
+def partitioning(
+    schema: Schema,
+) -> Partitioning: ...
 @overload
 def partitioning(
     schema: Schema,