From 616fc96ca86b4cb0fb4eba51b95cbdc50b537ff1 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 18 Jan 2024 12:46:09 +0100 Subject: [PATCH] Adding ability to calculate vector magnitude. --- supervision/geometry/core.py | 13 ++++++++++ .../{test_dataclasses.py => test_core.py} | 26 +++++++++++++++++++ 2 files changed, 39 insertions(+) rename test/geometry/{test_dataclasses.py => test_core.py} (59%) diff --git a/supervision/geometry/core.py b/supervision/geometry/core.py index 549d2c1a..468f43db 100644 --- a/supervision/geometry/core.py +++ b/supervision/geometry/core.py @@ -2,6 +2,7 @@ from __future__ import annotations from dataclasses import dataclass from enum import Enum +from math import sqrt from typing import Tuple @@ -43,6 +44,18 @@ class Vector: start: Point end: Point + @property + def magnitude(self) -> float: + """ + Calculate the magnitude (length) of the vector. + + Returns: + float: The magnitude of the vector. + """ + dx = self.end.x - self.start.x + dy = self.end.y - self.start.y + return sqrt(dx ** 2 + dy ** 2) + def cross_product(self, point: Point) -> float: """ Calculate the 2D cross product (also known as the vector product or outer diff --git a/test/geometry/test_dataclasses.py b/test/geometry/test_core.py similarity index 59% rename from test/geometry/test_dataclasses.py rename to test/geometry/test_core.py index 7f7a450a..da22eed1 100644 --- a/test/geometry/test_dataclasses.py +++ b/test/geometry/test_core.py @@ -33,3 +33,29 @@ def test_vector_cross_product( ) -> None: result = vector.cross_product(point=point) assert result == expected_result + + +@pytest.mark.parametrize( + "vector, expected_result", + [ + (Vector(start=Point(x=0, y=0), end=Point(x=0, y=0)), 0.0), + (Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), 1.0), + (Vector(start=Point(x=0, y=1), end=Point(x=0, y=0)), 1.0), + (Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), 1.0), + (Vector(start=Point(x=0, y=0), end=Point(x=0, y=1)), 1.0), + (Vector(start=Point(x=-1, y=0), end=Point(x=0, y=0)), 1.0), + (Vector(start=Point(x=0, y=-1), end=Point(x=0, y=0)), 1.0), + (Vector(start=Point(x=0, y=0), end=Point(x=-1, y=0)), 1.0), + (Vector(start=Point(x=0, y=0), end=Point(x=0, y=-1)), 1.0), + (Vector(start=Point(x=0, y=0), end=Point(x=3, y=4)), 5.0), + (Vector(start=Point(x=0, y=0), end=Point(x=-3, y=4)), 5.0), + (Vector(start=Point(x=0, y=0), end=Point(x=3, y=-4)), 5.0), + (Vector(start=Point(x=0, y=0), end=Point(x=-3, y=-4)), 5.0), + (Vector(start=Point(x=0, y=0), end=Point(x=4, y=3)), 5.0), + (Vector(start=Point(x=3, y=4), end=Point(x=0, y=0)), 5.0), + (Vector(start=Point(x=4, y=3), end=Point(x=0, y=0)), 5.0), + ] +) +def test_vector_magnitude(vector: Vector, expected_result: float) -> None: + result = vector.magnitude + assert result == expected_result