From a4ad4da4da637e3bd86e15d17361126c3e1e6072 Mon Sep 17 00:00:00 2001
From: Erik Friese <e.friese0@gmail.com>
Date: Mon, 1 Apr 2024 12:15:00 +0200
Subject: [PATCH] error messages improved

---
 src/betterproto_interop/error.rs | 10 +---------
 src/encode/error.rs              | 12 +++++++++---
 src/encode/mod.rs                |  2 +-
 src/well_known_types.rs          |  6 +++---
 4 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/src/betterproto_interop/error.rs b/src/betterproto_interop/error.rs
index 4879e97..42a00af 100644
--- a/src/betterproto_interop/error.rs
+++ b/src/betterproto_interop/error.rs
@@ -1,4 +1,4 @@
-use pyo3::{exceptions::PyRuntimeError, PyErr};
+use pyo3::PyErr;
 use thiserror::Error;
 
 #[derive(Error, Debug)]
@@ -13,14 +13,6 @@ pub enum InteropError {
     UnsupportedWrappedType(String),
     #[error("Given object is not a valid betterproto message.")]
     IncompleteMetadata,
-    #[error("Offset-naive datetime {0} is invalid for the current local timezone.")]
-    OffsetNaiveDateTimeDoesNotMap(chrono::NaiveDateTime),
 }
 
 pub type InteropResult<T> = Result<T, InteropError>;
-
-impl From<InteropError> for PyErr {
-    fn from(value: InteropError) -> Self {
-        PyRuntimeError::new_err(value.to_string())
-    }
-}
diff --git a/src/encode/error.rs b/src/encode/error.rs
index a253b82..3cd1d18 100644
--- a/src/encode/error.rs
+++ b/src/encode/error.rs
@@ -4,14 +4,16 @@ use thiserror::Error;
 
 #[derive(Error, Debug)]
 pub enum EncodeError {
-    #[error("Given object is not a valid betterproto message.")]
-    NoBetterprotoMessage(#[from] PyErr),
+    #[error(transparent)]
+    PythonInteropFailed(#[from] PyErr),
     #[error("Given object is not a valid betterproto message.")]
     DowncastFailed,
     #[error(transparent)]
     Interop(#[from] InteropError),
     #[error("Given object is not a valid betterproto message.")]
     ProstEncode(#[from] prost::EncodeError),
+    #[error("Offset-naive datetime {0} is invalid for the current local timezone.")]
+    OffsetNaiveDateTimeDoesNotMap(chrono::NaiveDateTime),
 }
 
 pub type EncodeResult<T> = Result<T, EncodeError>;
@@ -24,6 +26,10 @@ impl From<DowncastError<'_, '_>> for EncodeError {
 
 impl From<EncodeError> for PyErr {
     fn from(value: EncodeError) -> Self {
-        PyRuntimeError::new_err(value.to_string())
+        if let EncodeError::PythonInteropFailed(pyerr) = value {
+            pyerr
+        } else {
+            PyRuntimeError::new_err(value.to_string())
+        }
     }
 }
diff --git a/src/encode/mod.rs b/src/encode/mod.rs
index e02b981..b0b8c88 100644
--- a/src/encode/mod.rs
+++ b/src/encode/mod.rs
@@ -2,4 +2,4 @@ mod chunk;
 mod error;
 mod message;
 
-pub use self::{error::EncodeResult, message::MessageEncoder};
+pub use self::{error::EncodeError, error::EncodeResult, message::MessageEncoder};
diff --git a/src/well_known_types.rs b/src/well_known_types.rs
index f14dcf3..311780e 100644
--- a/src/well_known_types.rs
+++ b/src/well_known_types.rs
@@ -7,8 +7,8 @@ use pyo3::{
 };
 
 use crate::{
-    betterproto_interop::{InteropError, InteropResult},
     decode::{DecodeError, DecodeResult},
+    encode::{EncodeError, EncodeResult},
 };
 
 const NANOS_PER_SEC: u32 = 1_000_000_000;
@@ -200,7 +200,7 @@ impl Duration {
 }
 
 impl Timestamp {
-    fn try_from_any(ob: &Bound<PyAny>) -> InteropResult<Self> {
+    fn try_from_any(ob: &Bound<PyAny>) -> EncodeResult<Self> {
         // try to extract an offset-aware datetime object
         if let Ok(dt) = ob.extract::<chrono::DateTime<chrono::FixedOffset>>() {
             return Ok(dt.to_utc().into());
@@ -211,7 +211,7 @@ impl Timestamp {
         Ok(dt
             .and_local_timezone(chrono::Local)
             .single()
-            .ok_or(InteropError::OffsetNaiveDateTimeDoesNotMap(dt))?
+            .ok_or(EncodeError::OffsetNaiveDateTimeDoesNotMap(dt))?
             .to_utc()
             .into())
     }