Extend `stub.py` to accept external typehinting (#1102)

2 files changed, 58 insertions, 0 deletions

diff --git a/candle-pyo3/_additional_typing/README.md b/candle-pyo3/_additional_typing/README.md
new file mode 100644
index 00000000..ab5074e0
--- /dev/null
+++ b/candle-pyo3/_additional_typing/README.md
@@ -0,0 +1,3 @@
+This python module contains external typehinting for certain `candle` classes. This is only necessary for `magic` methodes e.g. `__add__` as their text signature cant be set via pyo3.
+
+The classes in this module will be parsed by the `stub.py` script and interleafed with the signatures of the actual pyo3 `candle.candle` module.
\ No newline at end of file
diff --git a/candle-pyo3/_additional_typing/__init__.py b/candle-pyo3/_additional_typing/__init__.py
new file mode 100644
index 00000000..0d0eec90
--- /dev/null
+++ b/candle-pyo3/_additional_typing/__init__.py
@@ -0,0 +1,55 @@
+from typing import Union, Sequence
+
+
+class Tensor:
+    """
+    This contains the type hints for the magic methodes of the `candle.Tensor` class.
+    """
+
+    def __add__(self, rhs: Union["Tensor", "Scalar"]) -> "Tensor":
+        """
+        Add a scalar to a tensor or two tensors together.
+        """
+        pass
+
+    def __radd__(self, rhs: Union["Tensor", "Scalar"]) -> "Tensor":
+        """
+        Add a scalar to a tensor or two tensors together.
+        """
+        pass
+
+    def __sub__(self, rhs: Union["Tensor", "Scalar"]) -> "Tensor":
+        """
+        Subtract a scalar from a tensor or one tensor from another.
+        """
+        pass
+
+    def __truediv__(self, rhs: Union["Tensor", "Scalar"]) -> "Tensor":
+        """
+        Divide a tensor by a scalar or one tensor by another.
+        """
+        pass
+
+    def __mul__(self, rhs: Union["Tensor", "Scalar"]) -> "Tensor":
+        """
+        Multiply a tensor by a scalar or one tensor by another.
+        """
+        pass
+
+    def __rmul__(self, rhs: Union["Tensor", "Scalar"]) -> "Tensor":
+        """
+        Multiply a tensor by a scalar or one tensor by another.
+        """
+        pass
+
+    def __richcmp__(self, rhs: Union["Tensor", "Scalar"], op) -> "Tensor":
+        """
+        Compare a tensor with a scalar or one tensor with another.
+        """
+        pass
+
+    def __getitem__(self, index: Union["Index", "Tensor", Sequence["Index"]]) -> "Tensor":
+        """
+        Return a slice of a tensor.
+        """
+        pass