From bb1077d45f19d543c15b836afc872cc54c3fb973 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Mon, 4 Jul 2022 08:47:52 +0200
Subject: [PATCH] Add exceptions

---
 src/great_ai/great_ai/__init__.py                            | 5 +++++
 src/great_ai/great_ai/exceptions/__init__.py                 | 1 +
 .../great_ai/exceptions/wrong_decorator_order_error.py       | 2 ++
 .../great_ai/helper/assert_function_is_not_finalised.py      | 5 ++---
 src/great_ai/great_ai/parameters/parameter.py                | 3 ++-
 5 files changed, 12 insertions(+), 4 deletions(-)
 create mode 100644 src/great_ai/great_ai/exceptions/wrong_decorator_order_error.py

diff --git a/src/great_ai/great_ai/__init__.py b/src/great_ai/great_ai/__init__.py
index ce37e74..f2014ca 100644
--- a/src/great_ai/great_ai/__init__.py
+++ b/src/great_ai/great_ai/__init__.py
@@ -1,5 +1,10 @@
 from .context import configure
 from .deploy import GreatAI
+from .exceptions import (
+    ArgumentValidationError,
+    MissingArgumentError,
+    WrongDecoratorOrderError,
+)
 from .models import save_model, use_model
 from .output_models import (
     ClassificationOutput,
diff --git a/src/great_ai/great_ai/exceptions/__init__.py b/src/great_ai/great_ai/exceptions/__init__.py
index 037b1aa..364e5ee 100644
--- a/src/great_ai/great_ai/exceptions/__init__.py
+++ b/src/great_ai/great_ai/exceptions/__init__.py
@@ -1,2 +1,3 @@
 from .argument_validation_error import ArgumentValidationError
 from .missing_argument_error import MissingArgumentError
+from .wrong_decorator_order_error import WrongDecoratorOrderError
diff --git a/src/great_ai/great_ai/exceptions/wrong_decorator_order_error.py b/src/great_ai/great_ai/exceptions/wrong_decorator_order_error.py
new file mode 100644
index 0000000..ca9864c
--- /dev/null
+++ b/src/great_ai/great_ai/exceptions/wrong_decorator_order_error.py
@@ -0,0 +1,2 @@
+class WrongDecoratorOrderError(Exception):
+    pass
diff --git a/src/great_ai/great_ai/helper/assert_function_is_not_finalised.py b/src/great_ai/great_ai/helper/assert_function_is_not_finalised.py
index b2110c2..55169ab 100644
--- a/src/great_ai/great_ai/helper/assert_function_is_not_finalised.py
+++ b/src/great_ai/great_ai/helper/assert_function_is_not_finalised.py
@@ -1,6 +1,6 @@
 from typing import Any, Callable
 
-from ..context import get_context
+from ..exceptions import WrongDecoratorOrderError
 from .get_function_metadata_store import get_function_metadata_store
 
 
@@ -11,5 +11,4 @@ def assert_function_is_not_finalised(func: Callable[..., Any]) -> None:
     )
 
     if get_function_metadata_store(func).is_finalised:
-        get_context().logger.error(error_message)
-        exit(-1)
+        raise WrongDecoratorOrderError(error_message)
diff --git a/src/great_ai/great_ai/parameters/parameter.py b/src/great_ai/great_ai/parameters/parameter.py
index 40e64eb..1d8e4b5 100644
--- a/src/great_ai/great_ai/parameters/parameter.py
+++ b/src/great_ai/great_ai/parameters/parameter.py
@@ -8,13 +8,14 @@ from ..helper import get_arguments, get_function_metadata_store
 from ..helper.assert_function_is_not_finalised import assert_function_is_not_finalised
 from ..tracing.tracing_context import TracingContext
 
+T = TypeVar("T")
 F = TypeVar("F", bound=Callable[..., Any])
 
 
 def parameter(
     parameter_name: str,
     *,
-    validator: Callable[[Any], bool] = lambda _: True,
+    validator: Callable[[T], bool] = lambda _: True,
     disable_logging: bool = False,
 ) -> Callable[[F], F]:
     def decorator(func: F) -> F: