From 61249f0632e3e85c751dcbca8c8ddcde6b825349 Mon Sep 17 00:00:00 2001
From: Jonas Loos <33965649+JonasLoos@users.noreply.github.com>
Date: Wed, 10 Apr 2024 19:39:56 +0200
Subject: [PATCH] add `StreamingDownloadManager` to
 `DatasetBuilder._split_generators` type annotation

---
 src/datasets/builder.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/datasets/builder.py b/src/datasets/builder.py
index 5a94b79e173..85f7d88b7db 100644
--- a/src/datasets/builder.py
+++ b/src/datasets/builder.py
@@ -1446,7 +1446,7 @@ def _download_post_processing_resources(
         return None
 
     @abc.abstractmethod
-    def _split_generators(self, dl_manager: DownloadManager):
+    def _split_generators(self, dl_manager: Union[DownloadManager, StreamingDownloadManager]):
         """Specify feature dictionary generators and dataset splits.
 
         This function returns a list of `SplitGenerator`s defining how to generate
@@ -1484,7 +1484,7 @@ def _split_generators(self, dl_manager: DownloadManager):
         distribute the relevant parts to each split with the `gen_kwargs` argument
 
         Args:
-            dl_manager (`DownloadManager`):
+            dl_manager (`Union[DownloadManager, StreamingDownloadManager]`):
                 Download manager to download the data
 
         Returns: