Add filtering logic for backtest saving and loading

Author: MDUYN

investing_algorithm_framework/app/analysis/backtest_utils.py

@@ -1,6 +1,6 @@
 import os
 from pathlib import Path
-from typing import List, Union
+from typing import List, Union, Callable
 from logging import getLogger
 from random import Random
 
@@ -12,7 +12,9 @@
 
 def save_backtests_to_directory(
     backtests: List[Backtest],
-    directory_path: Union[str, Path]
+    directory_path: Union[str, Path],
+    dir_name_generation_function: Callable[[Backtest], str] = None,
+    filter_function: Callable[[Backtest], bool] = None
 ) -> None:
     """
     Saves a list of Backtest objects to the specified directory.
@@ -21,6 +23,14 @@ def save_backtests_to_directory(
         backtests (List[Backtest]): List of Backtest objects to save.
         directory_path (str): Path to the directory where backtests
             will be saved.
+        dir_name_generation_function (Callable[[Backtest], str], optional):
+            A function that takes a Backtest object as input and returns
+            a string to be used as the directory name for that backtest.
+            If not provided, the backtest's metadata 'id' will be used.
+            Defaults to None.
+        filter_function (Callable[[Backtest], bool], optional): A function
+            that takes a Backtest object as input and returns True if the
+            backtest should be saved. Defaults to None.
 
     Returns:
         None
@@ -30,28 +40,40 @@ def save_backtests_to_directory(
         os.makedirs(directory_path)
 
     for backtest in backtests:
-        # Check if there is an ID in the backtest metadata
-        backtest_id = backtest.metadata.get('id')
 
-        if backtest_id is None:
+        if filter_function is not None:
+            if not filter_function(backtest):
+                continue
+
+        if dir_name_generation_function is not None:
+            dir_name = dir_name_generation_function(backtest)
+        else:
+            # Check if there is an ID in the backtest metadata
+            dir_name = backtest.metadata.get('id', None)
+
+        if dir_name is None:
             logger.warning(
-                "Backtest is missing 'id' in metadata. "
-                "Generating a random ID as name for backtest file."
+                "Backtest metadata does not contain an 'id' field. "
+                "Generating a random directory name."
             )
-            backtest_id = str(Random().randint(100000, 999999))
+            dir_name = str(Random().randint(100000, 999999))
 
-        backtest.save(os.path.join(directory_path, backtest_id))
+        backtest.save(os.path.join(directory_path, dir_name))
 
 
 def load_backtests_from_directory(
-    directory_path: Union[str, Path]
+    directory_path: Union[str, Path],
+    filter_function: Callable[[Backtest], bool] = None
 ) -> List[Backtest]:
     """
     Loads Backtest objects from the specified directory.
 
     Args:
         directory_path (str): Path to the directory from which backtests
             will be loaded.
+        filter_function (Callable[[Backtest], bool], optional): A function
+            that takes a Backtest object as input and returns True if the
+            backtest should be included in the result. Defaults to None.
 
     Returns:
         List[Backtest]: List of loaded Backtest objects.
@@ -71,6 +93,11 @@ def load_backtests_from_directory(
 
         try:
             backtest = Backtest.open(file_path)
+
+            if filter_function is not None:
+                if not filter_function(backtest):
+                    continue
+
             backtests.append(backtest)
         except Exception as e:
             logger.error(