add vec.__getitem__

include/typings/linalg.pyi

@@ -8,6 +8,7 @@ class vec2(_StructLike['vec2']):
     def __init__(self, x: float, y: float) -> None: ...
     def __add__(self, other: vec2) -> vec2: ...
     def __sub__(self, other: vec2) -> vec2: ...
+    def __getitem__(self, index: int) -> float: ...
 
     @overload
     def __mul__(self, other: float) -> vec2: ...
@@ -49,6 +50,7 @@ class vec3(_StructLike['vec3']):
     def __init__(self, x: float, y: float, z: float) -> None: ...
     def __add__(self, other: vec3) -> vec3: ...
     def __sub__(self, other: vec3) -> vec3: ...
+    def __getitem__(self, index: int) -> float: ...
 
     @overload
     def __mul__(self, other: float) -> vec3: ...
@@ -75,6 +77,7 @@ class vec4(_StructLike['vec4']):
     def __init__(self, x: float, y: float, z: float, w: float) -> None: ...
     def __add__(self, other: vec4) -> vec4: ...
     def __sub__(self, other: vec4) -> vec4: ...
+    def __getitem__(self, index: int) -> float: ...
 
     @overload
     def __mul__(self, other: float) -> vec4: ...

src/linalg.cpp

@@ -50,6 +50,15 @@ namespace pkpy{
             return VAR(self / other);                                                       \
         });
 
+#define BIND_VEC_GETITEM(D) \
+        vm->bind__getitem__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* obj, PyObject* index){ \
+            Vec##D& self = _CAST(Vec##D&, obj); \
+            i64 i = CAST(i64, index); \
+            if(i < 0 || i >= D) vm->IndexError("index out of range"); \
+            float* v = &self.x; \
+            return VAR(v[i]); \
+        });
+
 // https://github.com/Unity-Technologies/UnityCsReference/blob/master/Runtime/Export/Math/Vector2.cs#L289
 static Vec2 SmoothDamp(Vec2 current, Vec2 target, Vec2& currentVelocity, float smoothTime, float maxSpeed, float deltaTime)
 {
@@ -172,6 +181,7 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, Vec2& currentVelocity, float s
         BIND_VEC_FUNCTION_0(2, length_squared)
         BIND_VEC_FUNCTION_0(2, normalize)
         BIND_VEC_FUNCTION_0(2, normalize_)
+        BIND_VEC_GETITEM(2)
     }
 
     void Vec3::_register(VM* vm, PyObject* mod, PyObject* type){
@@ -206,6 +216,7 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, Vec2& currentVelocity, float s
         BIND_VEC_FUNCTION_0(3, length_squared)
         BIND_VEC_FUNCTION_0(3, normalize)
         BIND_VEC_FUNCTION_0(3, normalize_)
+        BIND_VEC_GETITEM(3)
     }
 
     void Vec4::_register(VM* vm, PyObject* mod, PyObject* type){
@@ -241,12 +252,14 @@ static Vec2 SmoothDamp(Vec2 current, Vec2 target, Vec2& currentVelocity, float s
         BIND_VEC_FUNCTION_0(4, length_squared)
         BIND_VEC_FUNCTION_0(4, normalize)
         BIND_VEC_FUNCTION_0(4, normalize_)
+        BIND_VEC_GETITEM(4)
     }
 
 #undef BIND_VEC_VEC_OP
 #undef BIND_VEC_MUL_OP
 #undef BIND_VEC_FUNCTION_0
 #undef BIND_VEC_FUNCTION_1
+#undef BIND_VEC_GETITEM
 
     void Mat3x3::_register(VM* vm, PyObject* mod, PyObject* type){
         PY_STRUCT_LIKE(Mat3x3)

tests/80_linalg.py

@@ -492,3 +492,8 @@ except IndexError:
 assert vec2(1, 2) * vec2(3, 4) == vec2(3, 8)
 assert vec3(1, 2, 3) * vec3(4, 5, 6) == vec3(4, 10, 18)
 assert vec4(1, 2, 3, 4) * vec4(5, 6, 7, 8) == vec4(5, 12, 21, 32)
+
+# test vec.__getitem__
+assert vec2(1, 2)[0] == 1 and vec2(1, 2)[1] == 2
+assert vec3(1, 2, 3)[0] == 1 and vec3(1, 2, 3)[1] == 2 and vec3(1, 2, 3)[2] == 3
+assert vec4(1, 2, 3, 4)[0] == 1 and vec4(1, 2, 3, 4)[1] == 2 and vec4(1, 2, 3, 4)[2] == 3 and vec4(1, 2, 3, 4)[3] == 4