diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 92ec7a2..b3dde92 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -28,7 +28,7 @@ repos: name: Run Type Checker types_or: [python, pyi, toml, yaml] pass_filenames: false - args: [quantio] + args: [quantio, test] - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.7.0 diff --git a/quantio/__init__.py b/quantio/__init__.py index 3775f58..0d5f6d8 100644 --- a/quantio/__init__.py +++ b/quantio/__init__.py @@ -2,6 +2,7 @@ from .exceptions import CanNotAddTypesError, CanNotSubtractTypesError from .quantities import Acceleration, Angle, Area, Length, Mass, Time, Velocity +from .vector import Vector __all__ = [ "Acceleration", @@ -11,6 +12,7 @@ "Mass", "Time", "Velocity", + "Vector", "CanNotAddTypesError", "CanNotSubtractTypesError", ] diff --git a/quantio/vector.py b/quantio/vector.py new file mode 100644 index 0000000..c76864b --- /dev/null +++ b/quantio/vector.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from typing import Generic, TypeVar + +import numpy as np + +T = TypeVar("T") + + +class Vector(Generic[T]): + """A vector of either quanity or numeric elements.""" + + _elements: np.array + + def __init__(self, elements: list | tuple | np.ndarray) -> None: + self._elements = np.array(elements) + + @classmethod + def __class_getitem__(cls, *_: object) -> type: + """Return this class for type hinting.""" + return cls diff --git a/test/test_quantities_base.py b/test/test_quantities_base.py index 05c4c67..258a92b 100644 --- a/test/test_quantities_base.py +++ b/test/test_quantities_base.py @@ -20,7 +20,7 @@ def test_add__false_class(): length = Length(meters=1) with pytest.raises(CanNotAddTypesError): - length += 1 + length += 1 # type: ignore def test_sub__success(): @@ -35,7 +35,7 @@ def test_sub__false_class(): length = Length(meters=1) with pytest.raises(CanNotSubtractTypesError): - length -= 1 + length -= 1 # type: ignore if __name__ == "__main__": diff --git a/test/test_vector.py b/test/test_vector.py new file mode 100644 index 0000000..32fac0e --- /dev/null +++ b/test/test_vector.py @@ -0,0 +1,19 @@ +import pytest + +import numpy as np + +from quantio import Vector, Length + + +def test_init(): + actual: Vector[Length] = Vector([Length.zero(), Length.zero()]) + assert np.all(actual._elements == np.array([Length.zero(), Length.zero()])) + + +def test_init_with_type_hint(): + actual = Vector[Length]([Length.zero(), Length.zero()]) + assert np.all(actual._elements == np.array([Length.zero(), Length.zero()])) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])