diff --git a/emmental/path.py b/emmental/path.py index 304d4be..67df2a0 100644 --- a/emmental/path.py +++ b/emmental/path.py @@ -1,9 +1,58 @@ # Copyright 2022 (c) Anna Schumaker. """Extra path handling for URIs.""" import pathlib +import threading import urllib +class ReaddirThread(threading.Thread): + """An object to manager asynchronous tree iteration.""" + + def __init__(self, directory: pathlib.Path): + """Initialize an IterTreeResult.""" + super().__init__() + self.root = directory + self._files = [] + self._lock = threading.Lock() + self._stop_event = threading.Event() + + def __read_directory(self, directory: pathlib.Path) -> None: + if self._stop_event.is_set(): + return + + for path in directory.iterdir(): + if path.is_dir(): + self.__read_directory(path) + else: + with self._lock: + self._files.append(path) + + def poll_result(self) -> list[pathlib.Path] | None: + """Poll for the result of the IterTreeThread.""" + with self._lock: + if self.is_alive() or len(self._files): + files = self._files + self._files = [] + return files + return None + + def run(self) -> None: + """Run the IterTreeThread.""" + self.__read_directory(self.root) + + def stop(self) -> None: + """Signal the thread to stop.""" + self._stop_event.set() + + +def readdir_async(directory: pathlib.Path) -> ReaddirThread: + """Iterate through a directory tree asynchronously.""" + if directory.is_dir(): + res = ReaddirThread(directory) + res.start() + return res + + def from_uri(uri: str) -> pathlib.Path: """Make a path from a uri.""" if parsed := urllib.parse.urlparse(uri): diff --git a/tests/test_path.py b/tests/test_path.py index bea53a2..eb5e91c 100644 --- a/tests/test_path.py +++ b/tests/test_path.py @@ -1,10 +1,79 @@ # Copyright 2022 (c) Anna Schumaker. """Tests our custom Path class.""" import pathlib +import threading import unittest +import unittest.mock import emmental.path +class TestReaddirAsync(unittest.TestCase): + """Test reading a directory tree asynchronously.""" + + def test_readdir_thread(self): + """Test the ReaddirThread object.""" + thread = emmental.path.ReaddirThread(pathlib.Path("/a/b/c")) + self.assertIsInstance(thread, threading.Thread) + self.assertIsInstance(thread._lock, type(threading.Lock())) + self.assertIsInstance(thread._stop_event, threading.Event) + self.assertListEqual(thread._files, []) + self.assertEqual(thread.root, pathlib.Path("/a/b/c")) + + self.assertFalse(thread.is_alive()) + + def test_readdir_thread_poll_result(self): + """Test the ReaddirThread poll_result() function.""" + thread = emmental.path.ReaddirThread(pathlib.Path("/a/b/c")) + with unittest.mock.patch.object(thread, "is_alive", + return_value=False): + self.assertIsNone(thread.poll_result()) + + with unittest.mock.patch.object(thread, "is_alive", + return_value=True): + self.assertListEqual(thread.poll_result(), []) + + files = [pathlib.Path("file1.txt"), pathlib.Path("file2.txt")] + thread._files = files + self.assertListEqual(thread.poll_result(), files) + self.assertListEqual(thread._files, []) + + def test_readdir_thread_read_directory(self): + """Test reading a directory.""" + file1 = unittest.mock.Mock(**{"is_dir.return_value": False}) + file2 = unittest.mock.Mock(**{"is_dir.return_value": False}) + dir1 = unittest.mock.Mock(**{"is_dir.return_value": True, + "iterdir.return_value": [file1]}) + dir2 = unittest.mock.Mock(**{"is_dir.return_value": True, + "iterdir.return_value": [dir1, file2]}) + + thread = emmental.path.ReaddirThread(pathlib.Path("/a/b/c")) + thread._ReaddirThread__read_directory(dir2) + self.assertListEqual(thread._files, [file1, file2]) + + def test_readdir_thread_stop(self): + """Test stopping a ReaddirThread.""" + thread = emmental.path.ReaddirThread(pathlib.Path("/a/b/c/")) + thread.stop() + self.assertTrue(thread._stop_event.is_set()) + + dir = unittest.mock.Mock() + dir.iterdir = unittest.mock.Mock(return_value=[]) + thread._ReaddirThread__read_directory(dir) + dir.iterdir.assert_not_called() + + @unittest.mock.patch("emmental.path.ReaddirThread.start") + def test_readdir_async(self, mock_start: unittest.mock.Mock): + """Test the readdir_async() function.""" + self.assertIsNone(emmental.path.readdir_async(pathlib.Path("/a/b/"))) + + dir = pathlib.Path("/a/b/c") + with unittest.mock.patch.object(pathlib.Path, "is_dir", + return_value=True): + thread = emmental.path.readdir_async(dir) + self.assertIsInstance(thread, emmental.path.ReaddirThread) + mock_start.assert_called() + + class TestPath(unittest.TestCase): """Test our path module."""