type annotations

Author: blarson-b10

lib/bindings/python/rust/http.rs

@@ -177,7 +177,7 @@ where
     Resp: Data + for<'de> Deserialize<'de>,
 {
     async fn generate(&self, request: SingleIn<Req>) -> Result<ManyOut<Annotated<Resp>>, Error> {
-        match self.0 .0.generate_in_parts(request).await {
+        match self.0.0.generate_in_parts(request).await {
             Ok((mut stream, context)) => {
                 let first_item = match futures::StreamExt::next(&mut stream).await {
                     Some(item) => item,

lib/bindings/python/rust/lib.rs

@@ -377,6 +377,12 @@ impl DistributedRuntime {
         self.inner.runtime().shutdown();
     }
 
+    fn child_token(&self) -> CancellationToken {
+        CancellationToken {
+            inner: self.inner.runtime().child_token(),
+        }
+    }
+
     fn event_loop(&self) -> PyObject {
         self.event_loop.clone()
     }

lib/bindings/python/src/dynamo/_core.pyi

@@ -52,6 +52,30 @@ class DistributedRuntime:
         Shutdown the runtime by triggering the cancellation token
         """
         ...
+
+    def child_token(self) -> CancellationToken:
+        """
+        Get a child cancellation token from the runtime
+        """
+        ...
+
+class CancellationToken:
+    """
+    A cancellation token for coordinating shutdown across components
+    """
+
+    def cancel(self) -> None:
+        """
+        Cancel the token
+        """
+        ...
+
+    async def cancelled(self) -> None:
+        """
+        Wait for the token to be cancelled
+        """
+        ...
+
 class EtcdClient:
     """
     Etcd is used for discovery in the DistributedRuntime