diff --git a/include/typings/linalg.pyi b/include/typings/linalg.pyi index 08353987..dc7f20c2 100644 --- a/include/typings/linalg.pyi +++ b/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: ... diff --git a/src/linalg.cpp b/src/linalg.cpp index 851a1e59..f1b104fb 100644 --- a/src/linalg.cpp +++ b/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) diff --git a/tests/80_linalg.py b/tests/80_linalg.py index 2314b11b..a3d69577 100644 --- a/tests/80_linalg.py +++ b/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