From 19fab0179ed8dbdb8b483cf42a5c9caed6353761 Mon Sep 17 00:00:00 2001
From: WyattBlue <wyattblue@auto-editor.com>
Date: Sun, 7 Apr 2024 13:47:10 -0400
Subject: [PATCH] Add type stubs for datasets.py

---
 av/audio/layout.pyi       | 1 +
 av/datasets.py            | 9 +++++----
 tests/test_audiolayout.py | 2 +-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/av/audio/layout.pyi b/av/audio/layout.pyi
index 338532e1a..a59398fcc 100644
--- a/av/audio/layout.pyi
+++ b/av/audio/layout.pyi
@@ -5,6 +5,7 @@ class AudioLayout:
     layout: int
     nb_channels: int
     channels: tuple[AudioChannel, ...]
+    def __init__(self, layout: int | str | AudioLayout): ...
 
 class AudioChannel:
     name: str
diff --git a/av/datasets.py b/av/datasets.py
index 15ffe9643..5954a9c98 100644
--- a/av/datasets.py
+++ b/av/datasets.py
@@ -2,12 +2,13 @@
 import logging
 import os
 import sys
+from typing import Iterator
 from urllib.request import urlopen
 
 log = logging.getLogger(__name__)
 
 
-def iter_data_dirs(check_writable=False):
+def iter_data_dirs(check_writable: bool = False) -> Iterator[str]:
     try:
         yield os.environ["PYAV_TESTDATA_DIR"]
     except KeyError:
@@ -42,7 +43,7 @@ def iter_data_dirs(check_writable=False):
     yield os.path.join(os.path.expanduser("~"), ".pyav", "datasets")
 
 
-def cached_download(url, name):
+def cached_download(url: str, name: str) -> str:
     """Download the data at a URL, and cache it under the given name.
 
     The file is stored under `pyav/test` with the given name in the directory
@@ -96,7 +97,7 @@ def cached_download(url, name):
     return path
 
 
-def fate(name):
+def fate(name: str) -> str:
     """Download and return a path to a sample from the FFmpeg test suite.
 
     Data is handled by :func:`cached_download`.
@@ -110,7 +111,7 @@ def fate(name):
     )
 
 
-def curated(name):
+def curated(name: str) -> str:
     """Download and return a path to a sample that is curated by the PyAV developers.
 
     Data is handled by :func:`cached_download`.
diff --git a/tests/test_audiolayout.py b/tests/test_audiolayout.py
index 0d13a2a18..9cecac19a 100644
--- a/tests/test_audiolayout.py
+++ b/tests/test_audiolayout.py
@@ -8,7 +8,7 @@ def test_stereo_properties(self):
         layout = AudioLayout("stereo")
         self._test_stereo(layout)
 
-    def test_2channel_properties(self):
+    def test_2channel_properties(self) -> None:
         layout = AudioLayout(2)
         self._test_stereo(layout)