Fix stateful processing using direct runner with type checks enabled

Closes apache#27167

sdks/python/apache_beam/transforms/util_test.py

@@ -39,6 +39,7 @@
 from apache_beam.metrics import MetricsFilter
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.options.pipeline_options import StandardOptions
+from apache_beam.options.pipeline_options import TypeOptions
 from apache_beam.portability import common_urns
 from apache_beam.portability.api import beam_runner_api_pb2
 from apache_beam.pvalue import AsList
@@ -1041,6 +1042,14 @@ def test_output_typehints(self):
               ShardedKeyType[typehints.Tuple[int, int]],  # type: ignore[misc]
               typehints.Iterable[str]])
 
+  def test_with_type_hints(self):
+    options = PipelineOptions()
+    options.view_as(TypeOptions).runtime_type_check = True
+    with TestPipeline(options=options) as pipeline:
+      pipeline \
+        | beam.Create(GroupIntoBatchesTest._create_test_data()) \
+        | util.GroupIntoBatches(GroupIntoBatchesTest.BATCH_SIZE)
+
   def _test_runner_api_round_trip(self, transform, urn):
     context = pipeline_context.PipelineContext()
     proto = transform.to_runner_api(context)

sdks/python/apache_beam/typehints/typecheck.py

@@ -48,6 +48,10 @@ class AbstractDoFnWrapper(DoFn):
   def __init__(self, dofn):
     super().__init__()
     self.dofn = dofn
+    if hasattr(dofn, 'on_window_timer'):
+      self.on_window_timer = dofn.on_window_timer
+    if hasattr(dofn, 'on_buffering_timer'):
+      self.on_buffering_timer = dofn.on_buffering_timer
 
   def _inspect_start_bundle(self):
     return self.dofn.get_function_arguments('start_bundle')