Implement module directory iteration.

Author: artob

Cargo.lock

@@ -21,7 +21,9 @@ asimov-config = { version = "25.1", default-features = false, features = [] }
 asimov-core = { version = "25.1", default-features = false, features = [] }
 asimov-directory = { version = "25.1", default-features = false, features = ["fs"] }
 derive_more = { version = "2", default-features = false, features = ["display"] }
-pyo3 = { version = "0.27", default-features = false, features = ["extension-module", "macros"] }
+pyo3 = { version = "0.28", default-features = false, features = ["extension-module", "macros"] }
+pyo3-async-runtimes = { version = "0.28", features = ["tokio-runtime"] }
+tokio = { version = "1", default-features = false, features = ["full"] }
 
 [patch.crates-io]
 asimov-config = { git = "https://github.com/asimov-platform/asimov.rs.git" }

Cargo.toml

@@ -31,9 +31,20 @@ class ModuleDirectory:
         pass
     def is_enabled(self, module_name: str) -> bool:
         pass
+    def __iter__(self) -> ModuleNameIterator:
+        pass
+    def __aiter__(self) -> ModuleNameIterator:
+        pass
 
 class ModuleNameIterator:
-    pass
+    def __iter__(self) -> ModuleNameIterator:
+        pass
+    def __aiter__(self) -> ModuleNameIterator:
+        pass
+    def __next__(self) -> str:
+        pass
+    async def __anext__(self) -> str:
+        pass
 
 class ProgramDirectory:
     @staticmethod

src/asimov/sdk.pyi

@@ -1,11 +1,20 @@
 // This is free and unencumbered software released into the public domain.
 
+use super::ModuleNameIterator;
+use crate::util::wait_for_future;
 use asimov_directory::{ModuleDirectory as _, fs::ModuleDirectory as FsModuleDirectory};
 use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyString};
+use std::sync::Arc;
 
 /// A module directory stored on a file system (e.g., `$HOME/.asimov/modules/`).
 #[pyclass(frozen, module = "asimov", str = "{0}")]
-pub struct ModuleDirectory(pub(crate) FsModuleDirectory);
+pub struct ModuleDirectory(pub(crate) Arc<FsModuleDirectory>);
+
+impl From<FsModuleDirectory> for ModuleDirectory {
+    fn from(input: FsModuleDirectory) -> Self {
+        Self(Arc::new(input))
+    }
+}
 
 #[pymethods]
 impl ModuleDirectory {
@@ -19,7 +28,7 @@ impl ModuleDirectory {
                 "Failed to open module directory", // TODO
             ))?
         };
-        Ok(Self(result))
+        Ok(Self::from(result))
     }
 
     #[new]
@@ -29,7 +38,7 @@ impl ModuleDirectory {
                 "Failed to open module directory", // TODO
             ))?
         };
-        Ok(Self(result))
+        Ok(Self::from(result))
     }
 
     pub fn __repr__(slf: &Bound<'_, Self>) -> PyResult<String> {
@@ -46,4 +55,18 @@ impl ModuleDirectory {
     pub fn is_enabled(&self, module_name: String) -> bool {
         self.0.is_enabled(module_name)
     }
+
+    /// Returns an iterator over the installed modules.
+    fn __iter__(&self, py: Python) -> PyResult<ModuleNameIterator> {
+        self.__aiter__(py)
+    }
+
+    /// Returns an iterator over the installed modules.
+    fn __aiter__(&self, py: Python) -> PyResult<ModuleNameIterator> {
+        let inner = Arc::clone(&self.0);
+        wait_for_future(py, async move {
+            let it = inner.iter_installed().await.unwrap();
+            ModuleNameIterator::from(it)
+        })
+    }
 }

src/directory/module_directory.rs

@@ -1,11 +1,59 @@
 // This is free and unencumbered software released into the public domain.
 
-use asimov_directory::fs::ModuleDirectory as FsModuleDirectory;
-use pyo3::prelude::*;
+use crate::util::wait_for_future;
+use asimov_directory::{ModuleNameIterator as _, fs::ModuleNameIterator as FsModuleNameIterator};
+use pyo3::{
+    exceptions::{PyStopAsyncIteration, PyStopIteration},
+    prelude::*,
+};
+use std::sync::Arc;
+use tokio::sync::Mutex;
 
-/// TODO
-#[pyclass(frozen, module = "asimov", str = "{0}")]
-pub struct ModuleNameIterator(pub(crate) FsModuleDirectory);
+/// An iterator over module names in the module directory.
+#[pyclass(frozen, module = "asimov", str = "ModuleNameIterator")]
+pub struct ModuleNameIterator(pub(crate) Arc<Mutex<FsModuleNameIterator>>);
+
+impl From<FsModuleNameIterator> for ModuleNameIterator {
+    fn from(input: FsModuleNameIterator) -> Self {
+        Self(Arc::new(Mutex::new(input)))
+    }
+}
 
 #[pymethods]
-impl ModuleNameIterator {}
+impl ModuleNameIterator {
+    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    fn __next__(&self, py: Python) -> PyResult<String> {
+        let inner = Arc::clone(&self.0);
+        let next = async move {
+            let mut lock = inner.lock().await;
+            let item = lock.next().await;
+            drop(lock);
+            match item {
+                Some(name) => Ok(name),
+                None => Err(PyStopIteration::new_err("")),
+            }
+        };
+        wait_for_future(py, next).flatten()
+    }
+
+    fn __anext__<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        let inner = Arc::clone(&self.0);
+        let anext = async move {
+            let mut lock = inner.lock().await;
+            let item = lock.next().await;
+            drop(lock);
+            match item {
+                Some(name) => Ok(name),
+                None => Err(PyStopAsyncIteration::new_err("")),
+            }
+        };
+        pyo3_async_runtimes::tokio::future_into_py(py, anext)
+    }
+}

src/directory/module_iterators.rs

@@ -70,7 +70,7 @@ impl StateDirectory {
                 "Failed to open module directory", // TODO
             ))?
         };
-        Ok(ModuleDirectory(result))
+        Ok(ModuleDirectory::from(result))
     }
 
     /// Opens the program directory under this state directory.

src/directory/state_directory.rs

@@ -34,3 +34,5 @@ mod sdk {
 
 mod directory;
 use directory::*;
+
+mod util;

src/lib.rs

@@ -0,0 +1,45 @@
+// This is free and unencumbered software released into the public domain.
+
+use pyo3::{PyResult, Python};
+use std::{sync::OnceLock, time::Duration};
+use tokio::runtime::Runtime;
+
+pub fn runtime() -> &'static Runtime {
+    static RUNTIME: OnceLock<Runtime> = OnceLock::new();
+    static RUNTIME_PID: OnceLock<u32> = OnceLock::new();
+    let current_pid = std::process::id();
+    let runtime_pid = *RUNTIME_PID.get_or_init(|| current_pid);
+    if current_pid != runtime_pid {
+        panic!("Forked process detected, bailing out."); // TODO
+    }
+    RUNTIME.get_or_init(|| Runtime::new().expect("should create the Tokio runtime"))
+}
+
+pub fn wait_for_future<F>(py: Python, fut: F) -> PyResult<F::Output>
+where
+    F: Future + Send,
+    F::Output: Send,
+{
+    let runtime: &Runtime = &runtime();
+    const INTERVAL_CHECK_SIGNALS: Duration = Duration::from_millis(500);
+
+    py.run(cr"pass", None, None)?;
+    py.check_signals()?;
+
+    py.detach(|| {
+        runtime.block_on(async {
+            tokio::pin!(fut);
+            loop {
+                tokio::select! {
+                    res = &mut fut => break Ok(res),
+                    _ = tokio::time::sleep(INTERVAL_CHECK_SIGNALS) => {
+                        Python::attach(|py| {
+                                py.run(cr"pass", None, None)?;
+                                py.check_signals()
+                        })?;
+                    }
+                }
+            }
+        })
+    })
+}