test.py

from typing import Callable
import inspect


class Test:
	def __init__(self):
		self._callback = self._passthrough
	
	@staticmethod
	def _passthrough(data, headers, *args):
		return data, headers

	def with_callback(self, callback: Callable):
		params_len = len(inspect.signature(callback).parameters)
		if params_len == 1:
			self._callback = lambda message, headers: (callback(message), headers)
		elif params_len == 2:
			self._callback = callback
		return self

	def start(self, message, headers):
		result, headers = self._callback(message, headers)
		print(f"{result}")
		print(f"{headers}")

def keep_headers(message):
	return f"Modified {message}"

def modify_headers(message, headers):
	return f"Modified {message}", f"Modified {headers}"

print("No callback")
Test().start("Message 1", "Headers 1")

print("Keep headers")
Test().with_callback(keep_headers).start("Message 1", "Headers 1")

print("Modify headers")
Test().with_callback(modify_headers).start("Message 2", "Headers 2")