fmt

# Conflicts:
#	rust/exo_rs/exo_rs.pyi
#	rust/exo_rs/pyproject.toml
#	rust/exo_rs/src/networking.rs
#	uv.lock

flake.nix

@@ -110,7 +110,7 @@
               nixpkgs-fmt.enable = true;
               ruff-format = {
                 enable = true;
-                excludes = [ "rust/exo_rs/exo_rs.pyi" ];
+                excludes = [ "rust/exo_rs/python/**/*.pyi" ];
               };
               rustfmt = {
                 enable = true;

python/parts.nix

@@ -46,7 +46,6 @@ let
       exoOverlay = final: prev: {
         # Replace workspace exo_rs with Nix-built wheel.
         # Preserve passthru so mkVirtualEnv can resolve dependency groups.
-        # Copy .pyi stub + py.typed marker so basedpyright can find the types.
         exo-rs = pkgs.stdenv.mkDerivation {
           pname = "exo-rs";
           version = "0.1.0";
@@ -55,12 +54,6 @@ let
           nativeBuildInputs = [ final.pyprojectWheelHook ];
           dontStrip = true;
           passthru = prev.exo-rs.passthru or { };
-
-          postInstall = ''
-            local siteDir=$out/${final.python.sitePackages}/exo_rs
-            cp ${inputs.self}/rust/exo_rs/exo_rs.pyi $siteDir/
-            touch $siteDir/py.typed
-          '';
         };
       };
       buildSystemsOverlay = final: prev:

rust/exo_rs/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "exo_rs"
-version = "0.2.16"
+version = "0.2.18"
 description = "Add your description here"
 readme = "README.md"
 authors = [
@@ -19,8 +19,8 @@ dev = ["exo_rs", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
 
 [tool.maturin]
 #purelib = true
-#python-source = "python"
-module-name = "exo_rs"
+python-source = "python"
+module-name = "exo_rs._core"
 features = ["pyo3/extension-module", "pyo3/experimental-async"]
 
 [tool.pyo3-stub-gen]
@@ -30,3 +30,6 @@ generate-init-py = true
 log_cli = true
 log_cli_level = "INFO"
 asyncio_mode = "auto"
+
+[tool.uv]
+cache-keys = [{ file = "src/**/*.rs" }]

rust/exo_rs/python/.gitignore

@@ -0,0 +1,4 @@
+*.so
+*.pyd
+*.dll
+*.dylib

rust/exo_rs/python/exo_rs/__init__.py

@@ -0,0 +1,4 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: F401
+
+__all__ = []

rust/exo_rs/python/exo_rs/_core/__init__.pyi

@@ -5,6 +5,7 @@ import builtins
 import os
 import pathlib
 import typing
+
 __all__ = [
     "AllQueuesFullError",
     "FromSwarm",
@@ -25,24 +26,35 @@ class AllQueuesFullError(builtins.Exception):
 class FromSwarm:
     @typing.final
     class Connection(FromSwarm):
-        __match_args__ = ("peer_id", "connected",)
+        __match_args__ = (
+            "peer_id",
+            "connected",
+        )
         @property
         def peer_id(self) -> builtins.str: ...
         @property
         def connected(self) -> builtins.bool: ...
-        def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> FromSwarm.Connection: ...
-    
+        def __new__(
+            cls, peer_id: builtins.str, connected: builtins.bool
+        ) -> FromSwarm.Connection: ...
+
     @typing.final
     class Message(FromSwarm):
-        __match_args__ = ("origin", "topic", "data",)
+        __match_args__ = (
+            "origin",
+            "topic",
+            "data",
+        )
         @property
         def origin(self) -> builtins.str: ...
         @property
         def topic(self) -> builtins.str: ...
         @property
         def data(self) -> bytes: ...
-        def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> FromSwarm.Message: ...
-    
+        def __new__(
+            cls, origin: builtins.str, topic: builtins.str, data: bytes
+        ) -> FromSwarm.Message: ...
+
     ...
 
 @typing.final
@@ -77,24 +89,29 @@ class MessageTooLargeError(builtins.Exception):
 
 @typing.final
 class NetworkingHandle:
-    def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
+    def __new__(
+        cls,
+        identity: Keypair,
+        bootstrap_peers: typing.Sequence[builtins.str],
+        listen_port: builtins.int,
+    ) -> NetworkingHandle: ...
     def recv(self) -> typing.Awaitable[FromSwarm]: ...
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.
-        
+
         Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
         """
     async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Unsubscribes from a `GossipSub` topic.
-        
+
         Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
         """
     async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
         r"""
         Publishes a message with multiple topics to the `GossipSub` network.
-        
+
         If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
         """
 
@@ -108,28 +125,30 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
 class Pidfile:
     r"""
     A PID file protected with a lock.
-    
+
     An instance of `Pidfile` can be used to manage a PID file: create it,
     lock it, detect already running daemons. It is backed by [`pidfile`]
     functions of `libbsd`/`libutil` which use `flopen` to lock the PID
     file.
-    
+
     When a PID file is created, the process ID of the current process is
     *not* written there, making it possible to lock the PID file before
     forking and only write the ID of the forked process when it is ready.
-    
+
     The PID file is deleted automatically when the `Pidfile` comes out of
     the scope. To close the PID file without deleting it, for example, in
     the parent process of a forked daemon, call `close()`.
-    
+
     [`exit`]: https://doc.rust-lang.org/std/process/fn.exit.html
     [`pidfile`]: https://linux.die.net/man/3/pidfile
     [`daemon`(3)]: https://linux.die.net/man/3/daemon
     """
-    def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
+    def __new__(
+        cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int
+    ) -> Pidfile:
         r"""
         Creates a new PID file and locks it.
-        
+
         If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
         a PID of the already running process, or `None` if no PID has been written to
         the PID file yet.
@@ -137,13 +156,13 @@ class Pidfile:
     def write(self) -> None:
         r"""
         Writes the current process ID to the PID file.
-        
+
         The file is truncated before writing.
         """
     def as_raw_fd(self) -> builtins.int:
         r"""
         Extracts the raw file descriptor.
-        
+
         This function is typically used to **borrow** an owned file descriptor.
         When used in this way, this method does **not** pass ownership of the
         raw file descriptor to the caller, and the file descriptor is only
@@ -159,4 +178,3 @@ class Pidfile:
 class PidfileError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
-

rust/exo_rs/python/exo_rs/ident/__init__.py

@@ -0,0 +1,8 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: F401
+
+from exo_rs._core import Keypair
+
+__all__ = [
+    "Keypair",
+]

rust/exo_rs/python/exo_rs/networking/__init__.py

@@ -0,0 +1,18 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: F401
+
+from exo_rs._core import (
+    AllQueuesFullError,
+    FromSwarm,
+    MessageTooLargeError,
+    NetworkingHandle,
+    NoPeersSubscribedToTopicError,
+)
+
+__all__ = [
+    "AllQueuesFullError",
+    "FromSwarm",
+    "MessageTooLargeError",
+    "NetworkingHandle",
+    "NoPeersSubscribedToTopicError",
+]

rust/exo_rs/python/exo_rs/pidfile/__init__.py

@@ -0,0 +1,9 @@
+# This file is automatically generated by pyo3_stub_gen
+# ruff: noqa: F401
+
+from exo_rs._core import Pidfile, PidfileError
+
+__all__ = [
+    "Pidfile",
+    "PidfileError",
+]

rust/exo_rs/src/lib.rs

@@ -1,21 +1,14 @@
-//! TODO: crate documentation
-//!
-//! this is here as a placeholder documentation
-//!
-//!
+//! Python package for EXO Rust bindings.
+
+module_doc!("exo_rs", "Python package for EXO Rust bindings.");
 
 mod allow_threading;
-mod ident;
-mod networking;
-mod pidfile;
+pub mod ident;
+pub mod networking;
+pub mod pidfile;
 
-use crate::ident::PyKeypair;
-use crate::networking::networking_submodule;
-use crate::pidfile::pidfile_submodule;
-use pyo3::prelude::PyModule;
-use pyo3::types::PyModuleMethods;
-use pyo3::{Bound, PyResult, pyclass, pymodule};
-use pyo3_stub_gen::define_stub_info_gatherer;
+use pyo3::{pyclass, pymodule};
+use pyo3_stub_gen::{define_stub_info_gatherer, module_doc, reexport_module_members};
 
 /// Namespace for all the constants used by this crate.
 pub(crate) mod r#const {
@@ -150,28 +143,43 @@ pub(crate) mod ext {
     }
 }
 
-/// A Python module implemented in Rust. The name of this function must match
-/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
-/// import the module.
-#[pymodule(name = "exo_rs", gil_used = true)]
-fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    // install logger
-    pyo3_log::init();
-    let mut builder = tokio::runtime::Builder::new_multi_thread();
-    builder.enable_all();
-    pyo3_async_runtimes::tokio::init(builder);
+#[pymodule(name = "_core", gil_used = true)]
+mod py_exo_rs {
+    #[pymodule_export]
+    use super::ident::PyKeypair;
+    #[pymodule_export]
+    use super::networking::{
+        PyAllQueuesFullError, PyFromSwarm, PyMessageTooLargeError, PyNetworkingHandle,
+        PyNoPeersSubscribedToTopicError,
+    };
+    #[pymodule_export]
+    use super::pidfile::{PyPidfile, PyPidfileError};
+    use pyo3::{
+        PyResult,
+        prelude::{Bound, PyModule},
+    };
 
-    // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
-    //       work with maturin, where the types generate correctly, in the right folder, without
-    //       too many importing issues...
-    m.add_class::<PyKeypair>()?;
-    networking_submodule(m)?;
-    pidfile_submodule(m)?;
+    #[pymodule_init]
+    fn init(_m: &Bound<'_, PyModule>) -> PyResult<()> {
+        // install logger (TODO: change to tracing)
+        pyo3_log::init();
 
-    // top-level constructs
-    // TODO: ...
+        // create pyo3_async_runtimes
+        let mut builder = tokio::runtime::Builder::new_multi_thread();
+        builder.enable_all();
+        pyo3_async_runtimes::tokio::init(builder);
 
-    Ok(())
+        Ok(())
+    }
 }
 
+// make sure these re-exports match the #[pymodule_export] from above
+reexport_module_members!("exo_rs.ident" from "exo_rs._core";
+    "Keypair");
+reexport_module_members!("exo_rs.networking" from "exo_rs._core";
+    "AllQueuesFullError", "FromSwarm", "MessageTooLargeError", "NetworkingHandle",
+    "NoPeersSubscribedToTopicError");
+reexport_module_members!("exo_rs.pidfile" from "exo_rs._core";
+    "Pidfile", "PidfileError");
+
 define_stub_info_gatherer!(stub_info);

rust/exo_rs/src/networking.rs

@@ -5,20 +5,20 @@ use crate::r#const::MPSC_CHANNEL_SIZE;
 use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
 use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
 use crate::ident::PyKeypair;
-use crate::networking::exception::{
-    PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
-};
 use crate::pyclass;
 use futures_lite::{Stream, StreamExt as _};
 use libp2p::gossipsub::PublishError;
 use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
 use pyo3::exceptions::PyRuntimeError;
-use pyo3::prelude::{PyModule, PyModuleMethods as _};
 use pyo3::types::PyBytes;
 use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods};
 use tokio::sync::{Mutex, mpsc, oneshot};
 
+pub use exception::{
+    PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
+};
+
 mod exception {
     use pyo3::types::PyTuple;
     use pyo3::{exceptions::PyException, prelude::*};
@@ -129,7 +129,7 @@ mod exception {
 
 #[gen_stub_pyclass]
 #[pyclass(name = "NetworkingHandle")]
-struct PyNetworkingHandle {
+pub struct PyNetworkingHandle {
     // channels
     pub to_swarm: mpsc::Sender<ToSwarm>,
     pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
@@ -137,7 +137,7 @@ struct PyNetworkingHandle {
 
 #[gen_stub_pyclass_complex_enum]
 #[pyclass(name = "FromSwarm")]
-enum PyFromSwarm {
+pub enum PyFromSwarm {
     Connection {
         peer_id: String,
         connected: bool,
@@ -148,6 +148,7 @@ enum PyFromSwarm {
         data: Py<PyBytes>,
     },
 }
+
 impl From<FromSwarm> for PyFromSwarm {
     fn from(value: FromSwarm) -> Self {
         match value {
@@ -296,14 +297,3 @@ impl PyNetworkingHandle {
         Ok(())
     }
 }
-
-pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
-    m.add_class::<exception::PyAllQueuesFullError>()?;
-    m.add_class::<exception::PyMessageTooLargeError>()?;
-
-    m.add_class::<PyNetworkingHandle>()?;
-    m.add_class::<PyFromSwarm>()?;
-
-    Ok(())
-}

rust/exo_rs/src/pidfile.rs

@@ -1,6 +1,5 @@
 use pidfile_rs::{Pidfile, PidfileError};
 use pyo3::exceptions::PyException;
-use pyo3::prelude::{PyModule, PyModuleMethods};
 use pyo3::{Bound, PyErr, PyResult, Python, pyclass, pymethods};
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
 use std::fs;
@@ -119,10 +118,3 @@ impl PyPidfile {
         self.0 = None;
     }
 }
-
-pub fn pidfile_submodule(m: &Bound<PyModule>) -> PyResult<()> {
-    m.add_class::<PyPidfileError>()?;
-    m.add_class::<PyPidfile>()?;
-
-    Ok(())
-}

rust/exo_rs/tests/test_python.py

@@ -1,14 +1,15 @@
 import asyncio
+from contextlib import suppress
 
 import pytest
 from _pytest.capture import CaptureFixture
-from exo_rs import (
-    Keypair,
+from exo_rs.ident import Keypair
+from exo_rs.networking import (
+    FromSwarm,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
-    Pidfile,
-    FromSwarm,
 )
+from exo_rs.pidfile import Pidfile
 
 
 @pytest.mark.asyncio
@@ -16,16 +17,21 @@ async def test_sleep_on_multiple_items() -> None:
     print("PYTHON: starting handle")
     h = NetworkingHandle(Keypair.generate(), [], 0)
 
-    rt = asyncio.create_task(_await_recv(h))
+    recv_task = asyncio.create_task(_await_recv(h))
 
-    # sleep for 4 ticks
-    for i in range(4):
-        await asyncio.sleep(1)
+    try:
+        # sleep for 4 ticks
+        for _ in range(4):
+            await asyncio.sleep(1)
 
-        try:
-            await h.gossipsub_publish("topic", b"somehting or other")
-        except NoPeersSubscribedToTopicError as e:
-            print("caught it", e)
+            try:
+                await h.gossipsub_publish("topic", b"somehting or other")
+            except NoPeersSubscribedToTopicError as e:
+                print("caught it", e)
+    finally:
+        recv_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await recv_task
 
 
 def test_pidfile(capsys: CaptureFixture[str]):
@@ -46,4 +52,5 @@ async def _await_recv(h: NetworkingHandle):
 
 
 def scoped_lock_file():
-    a = Pidfile("/tmp/lock.pid", 0o0600)
+    lock_file = Pidfile("/tmp/lock.pid", 0o0600)
+    lock_file.close()

rust/parts.nix

@@ -32,8 +32,17 @@
             isRootCargoFile =
               (baseName == "Cargo.toml" || baseName == "Cargo.lock")
               && (builtins.dirOf path == toString inputs.self);
+            isExoRsPythonSource =
+              (lib.hasInfix "/rust/exo_rs/python/" path || lib.hasSuffix "/rust/exo_rs/python" path)
+              && (
+                type == "directory"
+                || lib.hasSuffix ".py" path
+                || lib.hasSuffix ".pyi" path
+                || baseName == "py.typed"
+              );
           in
           isRootCargoFile
+          || isExoRsPythonSource
           || (inRustDir && (craneLib.filterCargoSources path type || lib.hasSuffix ".toml" path || lib.hasSuffix ".md" path));
       };
 

src/exo/main.py

@@ -10,7 +10,7 @@ from typing import Self
 import anyio
 from anyio.lowlevel import checkpoint as anyio_checkpoint
 from daemon import DaemonContext  # pyright: ignore[reportMissingTypeStubs]
-from exo_rs import Pidfile, PidfileError
+from exo_rs.pidfile import Pidfile, PidfileError
 from loguru import logger
 from pydantic import PositiveInt
 

src/exo/routing/connection_message.py

@@ -1,4 +1,4 @@
-from exo_rs import FromSwarm
+from exo_rs.networking import FromSwarm
 
 from exo.shared.types.common import NodeId
 from exo.utils.pydantic_ext import FrozenModel

src/exo/routing/router.py

@@ -12,10 +12,10 @@ from anyio import (
     move_on_after,
     sleep_forever,
 )
-from exo_rs import (
+from exo_rs.ident import Keypair
+from exo_rs.networking import (
     AllQueuesFullError,
     FromSwarm,
-    Keypair,
     MessageTooLargeError,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,

src/exo/utils/tests/test_pidfile.py

@@ -8,14 +8,14 @@ import textwrap
 from pathlib import Path
 from typing import Final
 
-from exo_rs import Pidfile
+from exo_rs.pidfile import Pidfile
 
 _CHILD_ACQUIRE_PIDFILE_SCRIPT: Final = textwrap.dedent(
     """
     import sys
     from pathlib import Path
 
-    from exo_rs import Pidfile, PidfileError
+    from exo_rs.pidfile import Pidfile, PidfileError
 
     path = Path(sys.argv[1])
     try:

uv.lock

@@ -623,7 +623,7 @@ requires-dist = [
 
 [[package]]
 name = "exo-rs"
-version = "0.2.16"
+version = "0.2.18"
 source = { editable = "rust/exo_rs" }
 
 [package.dev-dependencies]