From 63501c3b7919d25d4a6039f6fd180e6fd0194cb5 Mon Sep 17 00:00:00 2001
From: gmega <giuliano.mega@gmail.com>
Date: Thu, 20 Feb 2025 09:42:23 -0300
Subject: [PATCH] fix: simplify and fix adaptation of models to log entries

---
 benchmarks/logging/logging.py            | 12 ++++++------
 benchmarks/logging/tests/test_logging.py | 10 ++++++++++
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/benchmarks/logging/logging.py b/benchmarks/logging/logging.py
index 864e93d..9da11a1 100644
--- a/benchmarks/logging/logging.py
+++ b/benchmarks/logging/logging.py
@@ -47,19 +47,19 @@ class LogEntry(SnakeCaseModel):
         """Adapts an existing Pydantic model to a LogEntry. This is useful for when you have a model
         that you want to log and later recover from logs using :class:`LogParser` or :class:`LogSplitter`."""
 
-        def adapt_instance(cls, data: BaseModel):
-            return cls.model_validate(data.model_dump())
+        def adapt_instance(klass, data: BaseModel):
+            return klass.model_validate(data.model_dump())
 
         def recover_instance(self):
             return model.model_validate(self.model_dump())
 
-        parents = [base for base in model.__bases__ if issubclass(base, BaseModel)]
-
         adapted = type(
             f"{model.__name__}LogEntry",
-            tuple([LogEntry] + parents),
+            (
+                LogEntry,
+                model,
+            ),
             {
-                "__annotations__": model.__annotations__,
                 "adapt_instance": classmethod(adapt_instance),
                 "recover_instance": recover_instance,
             },
diff --git a/benchmarks/logging/tests/test_logging.py b/benchmarks/logging/tests/test_logging.py
index 0e6b460..5899a1c 100644
--- a/benchmarks/logging/tests/test_logging.py
+++ b/benchmarks/logging/tests/test_logging.py
@@ -247,6 +247,9 @@ class AModel(BaseModel):
 class BModel(AModel):
     b: int
 
+    def hello(self):
+        return f"world {self.a} {self.b}"
+
 
 def test_should_log_attributes_for_superclasses_in_adapted_entries():
     BModelLogEntry = LogEntry.adapt(BModel)
@@ -256,3 +259,10 @@ def test_should_log_attributes_for_superclasses_in_adapted_entries():
         str(BModelLogEntry.adapt_instance(instance))
         == '>>{"a":1,"b":2,"entry_type":"b_model_log_entry"}'
     )
+
+
+def test_should_preserve_methods_in_adapted_entries():
+    BModelLogEntry = LogEntry.adapt(BModel)
+    instance = BModel(a=1, b=2)
+
+    assert BModelLogEntry.adapt_instance(instance).hello() == "world 1 2"