From c3b6d44c59ae08ceb9a9315ed0245dee44cf3dfd Mon Sep 17 00:00:00 2001
From: Laurens Valk <laurens@pybricks.com>
Date: Mon, 30 May 2022 10:17:47 +0200
Subject: [PATCH] pybricks.geometry: Add typing.

---
 src/pybricks/geometry.py  | 87 ++++++++++++++++++++++++++++++++-------
 src/pybricks/geometry.pyi | 29 -------------
 2 files changed, 71 insertions(+), 45 deletions(-)
 delete mode 100644 src/pybricks/geometry.pyi

diff --git a/src/pybricks/geometry.py b/src/pybricks/geometry.py
index 20c068e..49e4042 100644
--- a/src/pybricks/geometry.py
+++ b/src/pybricks/geometry.py
@@ -1,16 +1,58 @@
 # SPDX-License-Identifier: MIT
-# Copyright (c) 2018-2021 The Pybricks Authors
+# Copyright (c) 2018-2022 The Pybricks Authors
 
 """Core linear algebra functionality for orientation sensors and robotics."""
 
+from __future__ import annotations
+
+from typing import Tuple, Collection, overload
+
 
 class Matrix:
-    """Mathematical representation of a matrix. It supports common operations
-    such as matrix addition (``+``), subtraction (``-``),
-    and multiplication (``*``). A :class:`.Matrix` object is immutable."""
+    """Mathematical representation of a matrix. It supports
+    addition (``A + B``), subtraction (``A - B``),
+    and matrix multiplication (``A * B``) for matrices of compatible size.
 
-    def __init__(self, rows):
-        """
+    It also supports scalar multiplication (``c * A`` or ``A * c``)
+    and scalar division (``A / c``).
+
+    A :class:`.Matrix` object is immutable."""
+
+    def __add__(self, other) -> Matrix:
+        ...
+
+    def __iadd__(self, other) -> Matrix:
+        ...
+
+    def __sub__(self, other) -> Matrix:
+        ...
+
+    def __isub__(self, other) -> Matrix:
+        ...
+
+    def __mul__(self, other) -> Matrix:
+        ...
+
+    def __rmul__(self, other) -> Matrix:
+        ...
+
+    def __imul__(self, other) -> Matrix:
+        ...
+
+    def __truediv__(self, other) -> Matrix:
+        ...
+
+    def __itruediv__(self, other) -> Matrix:
+        ...
+
+    def __floordiv__(self, other) -> Matrix:
+        ...
+
+    def __ifloordiv__(self, other) -> Matrix:
+        ...
+
+    def __init__(self, rows: Collection[Collection[int]]):
+        """Matrix(rows)
 
         Arguments:
             rows (list): List of rows. Each row is itself a list of numbers.
@@ -18,31 +60,44 @@ class Matrix:
         """
 
     @property
-    def T(self):
+    def T(self) -> Matrix:
         """Returns a new :class:`.Matrix` that is the transpose of the
         original."""
         pass
 
     @property
-    def shape(self):
+    def shape(self) -> Tuple[int, int]:
         """Returns a tuple (``m``, ``n``),
         where ``m`` is the number of rows and ``n`` is the number of columns.
         """
         pass
 
 
-def vector(x, y, z=None):
-    """Convenience function to create a :class:`.Matrix` with the
-    shape (``3``, ``1``) or (``2``, ``1``).
+@overload
+def vector(x: float, y: float) -> Matrix:
+    ...
+
+
+@overload
+def vector(x: float, y: float, z: float) -> Matrix:
+    ...
+
+
+def vector(*args):
+    """
+    vector(x, y) -> Matrix
+    vector(x, y, z) -> Matrix
+
+    Convenience function to create a :class:`.Matrix` with the
+    shape (``2``, ``1``) or (``3``, ``1``).
 
     Arguments:
         x (float): x-coordinate of the vector.
         y (float): y-coordinate of the vector.
         z (float): z-coordinate of the vector (optional).
 
-
     Returns:
-        Matrix: A matrix with the shape of a column vector.
+        A matrix with the shape of a column vector.
     """
     pass
 
@@ -56,6 +111,6 @@ class Axis:
 
     """
 
-    X = vector(1, 0, 0)
-    Y = vector(0, 1, 0)
-    Z = vector(0, 0, 1)
+    X: Matrix = vector(1, 0, 0)
+    Y: Matrix = vector(0, 1, 0)
+    Z: Matrix = vector(0, 0, 1)
diff --git a/src/pybricks/geometry.pyi b/src/pybricks/geometry.pyi
deleted file mode 100644
index 673cda6..0000000
--- a/src/pybricks/geometry.pyi
+++ /dev/null
@@ -1,29 +0,0 @@
-# SPDX-License-Identifier: MIT
-# Copyright (c) 2020-2021 The Pybricks Authors
-
-from typing import Collection, Optional, Tuple
-
-class Matrix:
-    def __init__(self, rows: Collection[Collection[int]]): ...
-    @property
-    def T(self) -> Matrix: ...
-    @property
-    def shape(self) -> Tuple[int, int]: ...
-    def __add__(self, other) -> Matrix: ...
-    def __iadd__(self, other) -> Matrix: ...
-    def __sub__(self, other) -> Matrix: ...
-    def __isub__(self, other) -> Matrix: ...
-    def __mul__(self, other) -> Matrix: ...
-    def __rmul__(self, other) -> Matrix: ...
-    def __imul__(self, other) -> Matrix: ...
-    def __truediv__(self, other) -> Matrix: ...
-    def __itruediv__(self, other) -> Matrix: ...
-    def __floordiv__(self, other) -> Matrix: ...
-    def __ifloordiv__(self, other) -> Matrix: ...
-
-def vector(x: float, y: float, z: Optional[float] = None) -> Matrix: ...
-
-class Axis:
-    X: Matrix
-    Y: Matrix
-    Z: Matrix