From 6391c5c29b9c03a23cae491d81eb8ca653b5444c Mon Sep 17 00:00:00 2001
From: Charles Turner <52199577+charles-turner-1@users.noreply.github.com>
Date: Tue, 10 Dec 2024 14:04:44 +0800
Subject: [PATCH] 297 translator errors (#298)

* Wrapped translators with more helpful errors & decorator that should centralise handling that - but is broken

* WIP

* Typo

* Fixed moving failures into translator

* Added test

* Updated test coverage

* Cleaned interface up a touch
---
 src/access_nri_intake/catalog/translators.py | 34 +++++++++++++++++++-
 tests/test_translators.py                    | 33 +++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/src/access_nri_intake/catalog/translators.py b/src/access_nri_intake/catalog/translators.py
index 84f215c..ea2dc09 100644
--- a/src/access_nri_intake/catalog/translators.py
+++ b/src/access_nri_intake/catalog/translators.py
@@ -7,7 +7,7 @@
 """
 
 from dataclasses import dataclass
-from functools import partial
+from functools import partial, wraps
 from typing import Callable
 
 import pandas as pd
@@ -17,6 +17,9 @@
 from . import COLUMNS_WITH_ITERABLES
 from .utils import _to_tuple, tuplify_series
 
+# Note: important that when using @tuplify_series and @trace_failure decorators,
+# trace failure is the innermost decorator
+
 FREQUENCY_TRANSLATIONS = {
     "monthly-averaged-by-hour": "1hr",
     "monthly-averaged-by-day": "1hr",
@@ -36,6 +39,29 @@
 }
 
 
+def trace_failure(func: Callable) -> Callable:
+    """
+    Decorator that wraps a function and prints a message if it raises an exception
+    """
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        func_name = func.__name__
+        colname = func_name[1:].split("_")[0]
+        # Ensure the first argument is an instance of the class
+        if not isinstance(args[0], DefaultTranslator):
+            raise TypeError("Decorator can only be applied to class methods")
+
+        try:
+            return func(*args, **kwargs)
+        except KeyError as exc:
+            raise KeyError(
+                f"Unable to translate '{colname}' column with translator '{args[0].__class__.__name__}'"
+            ) from exc
+
+    return wrapper
+
+
 class TranslatorError(Exception):
     "Generic Exception for the Translator classes"
 
@@ -192,6 +218,7 @@ def set_dispatch(
         self._dispatch[core_colname] = func
         setattr(self._dispatch_keys, core_colname, input_name)
 
+    @trace_failure
     def _realm_translator(self) -> pd.Series:
         """
         Return realm, fixing a few issues
@@ -199,6 +226,7 @@ def _realm_translator(self) -> pd.Series:
         return _cmip_realm_translator(self.source.df[self._dispatch_keys.realm])
 
     @tuplify_series
+    @trace_failure
     def _model_translator(self) -> pd.Series:
         """
         Return model from dispatch_keys.model
@@ -206,6 +234,7 @@ def _model_translator(self) -> pd.Series:
         return self.source.df[self._dispatch_keys.model]
 
     @tuplify_series
+    @trace_failure
     def _frequency_translator(self) -> pd.Series:
         """
         Return frequency, fixing a few issues
@@ -215,6 +244,7 @@ def _frequency_translator(self) -> pd.Series:
         )
 
     @tuplify_series
+    @trace_failure
     def _variable_translator(self) -> pd.Series:
         """
         Return variable as a tuple
@@ -407,6 +437,7 @@ def __init__(self, source, columns):
         )
 
     @tuplify_series
+    @trace_failure
     def _model_translator(self):
         """
         Get the model from the path. This is a slightly hacky approach, using the
@@ -425,6 +456,7 @@ def _realm_translator(self):
         return self.source.df.apply(lambda x: ("none",), 1)
 
     @tuplify_series
+    @trace_failure
     def _frequency_translator(self):
         """
         Get the frequency from the path
diff --git a/tests/test_translators.py b/tests/test_translators.py
index 4462efb..e8bfbaa 100644
--- a/tests/test_translators.py
+++ b/tests/test_translators.py
@@ -19,6 +19,7 @@
     TranslatorError,
     _cmip_realm_translator,
     _to_tuple,
+    trace_failure,
     tuplify_series,
 )
 
@@ -361,3 +362,35 @@ def test_NarclimTranslator(test_data, groupby, n_entries):
     esmds.description = "description"
     df = NarclimTranslator(esmds, CORE_COLUMNS).translate(groupby)
     assert len(df) == n_entries
+
+
+def test_translator_failure(test_data):
+    esmds = intake.open_esm_datastore(test_data / "esm_datastore/narclim2-zz63.json")
+    esmds.name = "name"
+    esmds.description = "description"
+    translator = NarclimTranslator(esmds, CORE_COLUMNS)
+
+    default = DefaultTranslator(esmds, CORE_COLUMNS)
+
+    translator.set_dispatch(
+        input_name="dud_name",
+        core_colname="model",
+        func=default._model_translator,
+    )
+
+    with pytest.raises(KeyError) as excinfo:
+        translator.translate()
+
+    assert (
+        "Unable to translate 'model' column with translator 'DefaultTranslator'"
+        in str(excinfo.value)
+    )
+
+    @trace_failure
+    def _(x: int) -> int:
+        return x
+
+    with pytest.raises(TypeError) as excinfo:
+        _(1)
+
+    assert "Decorator can only be applied to class methods" in str(excinfo.value)