optional timing of extractors

Author: Aske-Rosted

src/graphnet/data/extractors/extractor.py

@@ -23,7 +23,9 @@ class Extractor(ABC, Logger):
     An extractor is used in conjunction with a specific `FileReader`.
     """
 
-    def __init__(self, extractor_name: str, exclude: list = [None]):
+    def __init__(
+        self, extractor_name: str, exclude: list = [None], timed: bool = False
+    ) -> None:
         """Construct Extractor.
 
         Args:
@@ -32,10 +34,12 @@ def __init__(self, extractor_name: str, exclude: list = [None]):
                             data, and to name tables to which this data is
                             saved. E.g. "mc_truth".
             exclude: List of keys to exclude from the extracted data.
+            timed: Whether to time the execution of the __call__ method.
         """
         # Member variable(s)
         self._extractor_name: str = extractor_name
         self._exclude = exclude
+        self._timed = timed
 
         # Base class constructor
         super().__init__(name=__name__, class_name=self.__class__.__name__)
@@ -72,4 +76,24 @@ def name(self) -> str:
     def __init_subclass__(cls) -> None:
         """Initialize subclass and apply the exclude decorator to __call__."""
         super().__init_subclass__()
-        cls.__call__ = cls.exclude(cls.__call__)  # type: ignore
+        if not cls._timed:
+            cls.__call__ = cls.exclude(cls.__call__)  # type: ignore
+        else:
+            import time
+
+            cls.__call__ = cls.exclude(cls.__call__)  # type: ignore
+            # wrap with time logging
+            original_call = cls.__call__
+
+            def timed_call(
+                cls: "Extractor", *args: Any
+            ) -> Union[dict, pd.DataFrame]:
+                start_time = time.time()
+                result = original_call(cls, *args)
+                end_time = time.time()
+                cls._logger.info(
+                    f"Extractor {cls.name} took {end_time - start_time:.4f} seconds."
+                )
+                return result
+
+            cls.__call__ = timed_call  # type: ignore