Add type hinting

1 files changed, 5 insertions, 4 deletions

diff --git a/test/test_infinite.py b/test/test_infinite.py
index aa85ae5..bca21b6 100644
--- a/test/test_infinite.py
+++ b/test/test_infinite.py
@@ -1,5 +1,6 @@
 import pytest
 import numpy as np
+from numpy.typing import NDArray
 
 import matplotlib.pyplot as plt
 
@@ -8,17 +9,17 @@ from schroedinger import (
 )
 
 
-def psi(x, n, a):
+def psi(x: NDArray[np.float64], n: int, a: float) -> NDArray[np.float64]:
     n += 1 # Index starting from 0
     x = -x + a / 2 # Reflect x and move to the left
     return np.sqrt(2 / a) * np.sin(n * np.pi * x / a)
 
 
-def energy(n, a, m):
-    return (n + 1)**2 * np.pi**2 / m / a**2 / 2.0
+def energy(n: int, a: float, mass: float) -> float:
+    return (n + 1)**2 * np.pi**2 / mass / a**2 / 2.0
 
 
-def test_infinite():
+def test_infinite() -> None:
     conf = Config('test/infinite.inp')
     potential, delta = build_potential(conf)