KCL Python: Allow custom views, multiple screenshots of KCL (#7918)

This will allow teams to take multiple screenshots from different angles easily.

Author: adamchalmers

rust/kcl-python-bindings/src/bridge.rs

@@ -0,0 +1,37 @@
+//! For creating equivalents to the core kcmc types that support Python.
+use kittycad_modeling_cmds as kcmc;
+use pyo3::pyclass;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
+#[pyclass]
+pub struct Point3d {
+    x: f32,
+    y: f32,
+    z: f32,
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
+#[pyclass]
+pub struct CameraLookAt {
+    vantage: Point3d,
+    center: Point3d,
+    up: Point3d,
+}
+
+impl From<CameraLookAt> for kcmc::DefaultCameraLookAt {
+    fn from(CameraLookAt { vantage, center, up }: CameraLookAt) -> Self {
+        Self {
+            vantage: vantage.into(),
+            center: center.into(),
+            up: up.into(),
+            sequence: None,
+        }
+    }
+}
+
+impl From<Point3d> for kcmc::shared::Point3d<f32> {
+    fn from(Point3d { x, y, z }: Point3d) -> Self {
+        Self { x, y, z }
+    }
+}

rust/kcl-python-bindings/src/lib.rs

@@ -4,12 +4,15 @@ use kcl_lib::{
     lint::{checks, Discovered},
     ExecutorContext, UnitLength,
 };
+use kittycad_modeling_cmds as kcmc;
 use pyo3::{
     prelude::PyModuleMethods, pyclass, pyfunction, pymethods, pymodule, types::PyModule, wrap_pyfunction, Bound, PyErr,
     PyResult,
 };
 use serde::{Deserialize, Serialize};
 
+mod bridge;
+
 fn tokio() -> &'static tokio::runtime::Runtime {
     use std::sync::OnceLock;
     static RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
@@ -415,6 +418,27 @@ async fn execute_and_snapshot(path: String, image_format: ImageFormat) -> PyResu
 /// Execute the kcl code and snapshot it in a specific format.
 #[pyfunction]
 async fn execute_code_and_snapshot(code: String, image_format: ImageFormat) -> PyResult<Vec<u8>> {
+    let mut snaps = execute_code_and_snapshot_at_views(code, image_format, Vec::new()).await?;
+    Ok(snaps.pop().unwrap())
+}
+
+#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
+#[pyclass]
+struct SnapshotOptions {
+    /// If none, will use isometric view.
+    camera: Option<bridge::CameraLookAt>,
+    padding: f32,
+}
+
+/// Execute the kcl code and snapshot it in a specific format.
+/// Returns one image for each camera angle you provide.
+/// If you don't provide any camera angles, a default head-on camera angle will be used.
+#[pyfunction]
+async fn execute_code_and_snapshot_at_views(
+    code: String,
+    image_format: ImageFormat,
+    snapshot_options: Vec<SnapshotOptions>,
+) -> PyResult<Vec<Vec<u8>>> {
     tokio()
         .spawn(async move {
             let program =
@@ -428,46 +452,72 @@ async fn execute_code_and_snapshot(code: String, image_format: ImageFormat) -> P
                 .await
                 .map_err(|err| into_miette(err, &code))?;
 
-            // Zoom to fit.
-            ctx.engine
-                .send_modeling_cmd(
-                    uuid::Uuid::new_v4(),
-                    kcl_lib::SourceRange::default(),
-                    &kittycad_modeling_cmds::ModelingCmd::ZoomToFit(kittycad_modeling_cmds::ZoomToFit {
-                        object_ids: Default::default(),
-                        padding: 0.1,
-                        animated: false,
-                    }),
-                )
-                .await?;
-
-            // Send a snapshot request to the engine.
-            let resp = ctx
-                .engine
-                .send_modeling_cmd(
-                    uuid::Uuid::new_v4(),
-                    kcl_lib::SourceRange::default(),
-                    &kittycad_modeling_cmds::ModelingCmd::TakeSnapshot(kittycad_modeling_cmds::TakeSnapshot {
-                        format: image_format.into(),
-                    }),
-                )
-                .await?;
-
-            let kittycad_modeling_cmds::websocket::OkWebSocketResponseData::Modeling {
-                modeling_response: kittycad_modeling_cmds::ok_response::OkModelingCmdResponse::TakeSnapshot(data),
-            } = resp
-            else {
-                return Err(pyo3::exceptions::PyException::new_err(format!(
-                    "Unexpected response from engine: {resp:?}"
-                )));
-            };
+            if snapshot_options.is_empty() {
+                let data_bytes = snapshot(&ctx, image_format, 0.1).await?;
+                return Ok(vec![data_bytes]);
+            }
 
-            Ok(data.contents.0)
+            let mut snaps = Vec::with_capacity(snapshot_options.len());
+            for pre_snap in snapshot_options {
+                if let Some(camera) = pre_snap.camera {
+                    let view_cmd = kcmc::DefaultCameraLookAt::from(camera);
+                    let view_cmd = kcmc::ModelingCmd::DefaultCameraLookAt(view_cmd);
+                    ctx.engine
+                        .send_modeling_cmd(uuid::Uuid::new_v4(), Default::default(), &view_cmd)
+                        .await?;
+                } else {
+                    let view_cmd = kcmc::ModelingCmd::ViewIsometric(kcmc::ViewIsometric { padding: 0.0 });
+                    ctx.engine
+                        .send_modeling_cmd(uuid::Uuid::new_v4(), Default::default(), &view_cmd)
+                        .await?;
+                }
+                let data_bytes = snapshot(&ctx, image_format, pre_snap.padding).await?;
+                snaps.push(data_bytes);
+            }
+            Ok(snaps)
         })
         .await
         .map_err(|err| pyo3::exceptions::PyException::new_err(err.to_string()))?
 }
 
+async fn snapshot(ctx: &ExecutorContext, image_format: ImageFormat, padding: f32) -> PyResult<Vec<u8>> {
+    // Zoom to fit.
+    ctx.engine
+        .send_modeling_cmd(
+            uuid::Uuid::new_v4(),
+            kcl_lib::SourceRange::default(),
+            &kittycad_modeling_cmds::ModelingCmd::ZoomToFit(kittycad_modeling_cmds::ZoomToFit {
+                object_ids: Default::default(),
+                padding,
+                animated: false,
+            }),
+        )
+        .await?;
+
+    // Send a snapshot request to the engine.
+    let resp = ctx
+        .engine
+        .send_modeling_cmd(
+            uuid::Uuid::new_v4(),
+            kcl_lib::SourceRange::default(),
+            &kittycad_modeling_cmds::ModelingCmd::TakeSnapshot(kittycad_modeling_cmds::TakeSnapshot {
+                format: image_format.into(),
+            }),
+        )
+        .await?;
+
+    let kittycad_modeling_cmds::websocket::OkWebSocketResponseData::Modeling {
+        modeling_response: kittycad_modeling_cmds::ok_response::OkModelingCmdResponse::TakeSnapshot(data),
+    } = resp
+    else {
+        return Err(pyo3::exceptions::PyException::new_err(format!(
+            "Unexpected response from engine: {resp:?}",
+        )));
+    };
+
+    Ok(data.contents.0)
+}
+
 /// Execute a kcl file and export it to a specific file format.
 #[pyfunction]
 async fn execute_and_export(path: String, export_format: FileExportFormat) -> PyResult<Vec<ExportFile>> {
@@ -617,6 +667,7 @@ fn kcl(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(mock_execute_code, m)?)?;
     m.add_function(wrap_pyfunction!(execute_and_snapshot, m)?)?;
     m.add_function(wrap_pyfunction!(execute_code_and_snapshot, m)?)?;
+    m.add_function(wrap_pyfunction!(execute_code_and_snapshot_at_views, m)?)?;
     m.add_function(wrap_pyfunction!(execute_and_export, m)?)?;
     m.add_function(wrap_pyfunction!(execute_code_and_export, m)?)?;
     m.add_function(wrap_pyfunction!(format, m)?)?;