Switched public interface to use pint Quantity to describe input

2 files changed, 21 insertions, 7 deletions

diff --git a/setup.py b/setup.py
index 3adf719..8ade850 100755
--- a/setup.py
+++ b/setup.py
@@ -5,6 +5,7 @@ from setuptools import setup
 setup(
 	name="szilagyi",
 	packages=["szilagyi"],
+	install_requires=["pint"],
 	extra_requires={"plots": ["matplotlib"]},
 	package_data={"": ["_dataset/*.csv"]},
 	include_package_data=True,
diff --git a/szilagyi/__init__.py b/szilagyi/__init__.py
index 3e91936..8d294d7 100644
--- a/szilagyi/__init__.py
+++ b/szilagyi/__init__.py
@@ -4,25 +4,38 @@ are continuous. The wind-speed at 850 hPa is ignored as the nomogram dictates it
 
 There is only one function provided, so the use should be straight-forward. For example::
 
+	>>> import pint
 	>>> import szilagyi
-	>>> round(szilagyi.calculate_swi(10.5, 24100), 3)
+	>>> u = pint.UnitRegistry()
+	>>> round(szilagyi.calculate_swi(10.5 * u.C, 24100 * u.ft), 3)
 	0.811
-	>>> round(szilagyi.calculate_swi(7, 12500), 3)
+	>>> round(szilagyi.calculate_swi(7 * u.C, 12500 * u.ft), 3)
 	-4.88
 """
 
+from pint import UnitRegistry
+from pint.quantity import Quantity
+
 from ._nomogram import calculate_swi as _calculate_swi
 from ._nomogram import MAX_TEMPERATURE_DIFFERENCE, MAX_CONVECTIVE_CLOUD_DEPTH
 
+_registry = UnitRegistry()
+
 
-def calculate_swi(temperature_difference, convective_cloud_depth):
+def calculate_swi(temperature_difference: Quantity, convective_cloud_depth: Quantity) -> float:
 	"""
-	Calculates the Szilagyi Waterspout Index for *temperature_difference* and *convective_cloud_depth*. These are
-	expected to be expressed in Celsius degrees and feet, respectively.
+	Calculates the Szilagyi Waterspout Index for *temperature_difference* and *convective_cloud_depth*.
 
-	Function will raise ValueError if one of the values is not in range of the original nomogram. The maximum values are
-	available as MAX_TEMPERATURE_DIFFERENCE and MAX_CONVECTIVE_CLOUD_DEPTH constants.
+	Function will raise ValueError if input values are not convertible to their expected units or if they are not in range
+	of the original nomogram. The maximum values are available as MAX_TEMPERATURE_DIFFERENCE and MAX_CONVECTIVE_CLOUD_DEPTH
+	constants.
 	"""
+	if not temperature_difference.is_compatible_with(_registry.C):
+		raise ValueError("temperature_difference must be convertible to Celsius degrees")
+	if not convective_cloud_depth.is_compatible_with(_registry.ft):
+		raise ValueError("convective_cloud_depth must be convertible to Feet")
+	temperature_difference = temperature_difference.m_as(_registry.C)
+	convective_cloud_depth = convective_cloud_depth.m_as(_registry.ft)
 	if temperature_difference < 0 or temperature_difference > MAX_TEMPERATURE_DIFFERENCE:
 		raise ValueError(f"temperature_difference must be within <0, {MAX_TEMPERATURE_DIFFERENCE}> range")
 	if convective_cloud_depth < 0 or convective_cloud_depth > MAX_CONVECTIVE_CLOUD_DEPTH: