Improve simple strategy, automate __display_name__

1 files changed, 20 insertions, 0 deletions

diff --git a/robusta_krr/utils/display_name.py b/robusta_krr/utils/display_name.py
new file mode 100644
index 0000000..3cc89ce
--- /dev/null
+++ b/robusta_krr/utils/display_name.py
@@ -0,0 +1,20 @@
+from typing import Callable, TypeVar
+
+_T = TypeVar("_T")
+
+
+def add_display_name(*, postfix: str) -> Callable[[type[_T]], type[_T]]:
+    """Add a decorator factory to add __display_name__ property to the class."""
+
+    def decorator(cls: type[_T]) -> type[_T]:
+        class DisplayNameProperty:
+            def __get__(self, instance, owner):
+                if owner.__name__.lower().endswith(postfix.lower()):
+                    return owner.__name__[: -len(postfix)]
+
+                return owner.__name__
+
+        cls.__display_name__ = DisplayNameProperty()
+        return cls
+
+    return decorator