From 38ef421442c0f57d6cba83952c2f7a34650e1bae Mon Sep 17 00:00:00 2001
From: Joey Tran <joey.tran@schrodinger.com>
Date: Fri, 11 Oct 2024 10:16:38 -0400
Subject: [PATCH] Update typecheck err msg

---
 .../apache_beam/transforms/ptransform.py      | 23 +++++++++++++------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/sdks/python/apache_beam/transforms/ptransform.py b/sdks/python/apache_beam/transforms/ptransform.py
index 8554ebce5dbd..065b1416712c 100644
--- a/sdks/python/apache_beam/transforms/ptransform.py
+++ b/sdks/python/apache_beam/transforms/ptransform.py
@@ -946,14 +946,23 @@ def element_type(side_input):
           continue
         if not typehints.is_consistent_with(bindings.get(arg, typehints.Any),
                                             hint):
+          #raise TypeCheckError(
+          #'Type hint violation for \'{label}\': requires {hint} but got '
+          #'{actual_type} for {arg}\nFull type hint:\n{debug_str}'.format(
+          #label=self.label,
+          #hint=hint,
+          #actual_type=bindings[arg],
+          #arg=arg,
+          #debug_str=type_hints.debug_str()))
+          transform_nest_level = self.label.count("/")
+          split_producer_label = pvalueish.producer.full_label.split("/")
+          producer_label = "/".join(
+              split_producer_label[:transform_nest_level + 1])
           raise TypeCheckError(
-              'Type hint violation for \'{label}\': requires {hint} but got '
-              '{actual_type} for {arg}\nFull type hint:\n{debug_str}'.format(
-                  label=self.label,
-                  hint=hint,
-                  actual_type=bindings[arg],
-                  arg=arg,
-                  debug_str=type_hints.debug_str()))
+              f"The transform '{self.label}' only accepts PCollections of type '{hint}' "
+              f"but was applied to a PCollection of type '{bindings[arg]}' (produced by the transform '{producer_label}'). "
+              "Please ensure the input PCollection contains elements of the correct type."
+          )
 
   def _process_argspec_fn(self):
     """Returns an argspec of the function actually consuming the data.