Swizzles.

Author: tychedelia

crates/processing_pyo3/src/math.rs

@@ -302,31 +302,26 @@ macro_rules! impl_py_vec {
                 let py = slf.py();
                 let chars: Vec<char> = name.chars().collect();
                 let len = chars.len();
-                if !(2..=4).contains(&len) {
-                    return Err(PyAttributeError::new_err(format!(
+                let attr_err = || {
+                    PyAttributeError::new_err(format!(
                         "'{}' object has no attribute '{}'",
                         $py_name, name
-                    )));
+                    ))
+                };
+                if !(2..=4).contains(&len) {
+                    return Err(attr_err());
                 }
                 let mut vals = [0.0f32; 4];
                 for (i, c) in chars.iter().enumerate() {
-                    let idx = match c {
-                        'x' => 0usize,
+                    let idx: usize = match c {
+                        'x' => 0,
                         'y' => 1,
                         'z' => 2,
                         'w' => 3,
-                        _ => {
-                            return Err(PyAttributeError::new_err(format!(
-                                "'{}' object has no attribute '{}'",
-                                $py_name, name
-                            )));
-                        }
+                        _ => return Err(attr_err()),
                     };
                     if idx >= $n {
-                        return Err(PyAttributeError::new_err(format!(
-                            "'{}' object has no attribute '{}'",
-                            $py_name, name
-                        )));
+                        return Err(attr_err());
                     }
                     vals[i] = slf.0[idx];
                 }
@@ -566,13 +561,6 @@ impl_py_vec!(PyVec3, "Vec3", 3, [(x, set_x, 0), (y, set_y, 1), (z, set_z, 2)], V
     fn truncate(&self) -> PyVec2 {
         PyVec2(self.0.truncate())
     }
-
-    #[getter]
-    fn xy(&self) -> PyVec2 { PyVec2(self.0.xy()) }
-    #[getter]
-    fn xz(&self) -> PyVec2 { PyVec2(self.0.xz()) }
-    #[getter]
-    fn yz(&self) -> PyVec2 { PyVec2(self.0.yz()) }
 });
 
 impl_py_vec!(
@@ -585,11 +573,6 @@ impl_py_vec!(
         fn truncate(&self) -> PyVec3 {
             PyVec3(self.0.truncate())
         }
-
-        #[getter]
-        fn xy(&self) -> PyVec2 { PyVec2(self.0.xy()) }
-        #[getter]
-        fn xyz(&self) -> PyVec3 { PyVec3(self.0.xyz()) }
     }
 );
 

tools/generate_stubs/src/main.rs

@@ -1,7 +1,72 @@
+use pyo3_introspection::model::{Argument, Arguments, Expr, Function, Module};
 use pyo3_introspection::{introspect_cdylib, module_stub_files};
 use std::path::{Path, PathBuf};
 use std::{env, fs};
 
+const SWIZZLE_CHARS: [char; 4] = ['x', 'y', 'z', 'w'];
+
+fn swizzle_props(dim: usize) -> Vec<Function> {
+    let chars = &SWIZZLE_CHARS[..dim];
+    let mut out = Vec::new();
+    for length in 2..=4 {
+        let count = chars.len().pow(length as u32);
+        for n in 0..count {
+            let mut name = String::with_capacity(length);
+            let mut idx = n;
+            for _ in 0..length {
+                name.push(chars[idx % chars.len()]);
+                idx /= chars.len();
+            }
+            out.push(Function {
+                name,
+                decorators: vec![Expr::Name {
+                    id: "property".into(),
+                }],
+                arguments: Arguments {
+                    positional_only_arguments: vec![Argument {
+                        name: "self".into(),
+                        default_value: None,
+                        annotation: None,
+                    }],
+                    arguments: vec![],
+                    vararg: None,
+                    keyword_only_arguments: vec![],
+                    kwarg: None,
+                },
+                returns: Some(Expr::Attribute {
+                    value: Box::new(Expr::Name {
+                        id: "mewnala.math".into(),
+                    }),
+                    attr: format!("Vec{length}"),
+                }),
+                is_async: false,
+                docstring: None,
+            });
+        }
+    }
+    out
+}
+
+fn inject_swizzles(module: &mut Module) {
+    for math in module.modules.iter_mut().filter(|m| m.name == "math") {
+        for cls in math.classes.iter_mut() {
+            let dim = match cls.name.as_str() {
+                "Vec2" => 2,
+                "Vec3" => 3,
+                "Vec4" => 4,
+                _ => continue,
+            };
+            let existing: std::collections::HashSet<String> =
+                cls.methods.iter().map(|m| m.name.clone()).collect();
+            cls.methods.extend(
+                swizzle_props(dim)
+                    .into_iter()
+                    .filter(|m| !existing.contains(&m.name)),
+            );
+        }
+    }
+}
+
 fn workspace_root() -> &'static Path {
     Path::new(env!("CARGO_MANIFEST_DIR"))
         .parent()
@@ -45,6 +110,7 @@ fn main() {
         introspect_cdylib(&cdylib_path, "mewnala").expect("Failed to introspect cdylib");
 
     module.incomplete = false;
+    inject_swizzles(&mut module);
 
     let stubs = module_stub_files(&module);