Making timed decorator compatible with classes

tests/test_usage.py

@@ -164,6 +164,17 @@ def wait(self, x):
         key = next(iter(out.keys()))
         self.assertIn(ClassA.__name__, key)
 
+    def test_class_name(self):
+        out = {}
+
+        class ClassA:
+            def __call__(self, x):
+                sleep(x)
+
+        timed(out=out)(ClassA())(0.1)
+        key = next(iter(out.keys()))
+        self.assertTrue(key.endswith("<locals>.ClassA"))
+
     def test_return_time(self):
         seconds = 0.1
 

timed_decorator/utils.py

@@ -167,12 +167,18 @@ def update_timing_dict(out: dict, key: str, elapsed: int, own_time: int):
 def get_fn_name(fn):
     if hasattr(fn, "name"):  # torch.jit.ScriptFunction uses name instead of __name__
         return fn.name
-    return fn.__name__
+    if hasattr(fn, "__name__"):
+        return fn.__name__
+    if hasattr(fn, "__class__"):
+        return fn.__class__.__name__
+    return type(fn).__name__
 
 
 def get_fn_qualname(fn):
-    if hasattr(
-        fn, "qualified_name"
-    ):  # torch.jit.ScriptFunction uses qualified_name instead of __qualname__
+    if hasattr(fn, "qualified_name"):  # torch.jit.ScriptFunction uses qualified_name instead of __qualname__
         return fn.qualified_name
-    return fn.__qualname__
+    if hasattr(fn, "__qualname__"):
+        return fn.__qualname__
+    if hasattr(fn, "__class__"):
+        return fn.__class__.__qualname__
+    return type(fn).__qualname__