From 72f56330617217180c67e817a29010ce935a3980 Mon Sep 17 00:00:00 2001
From: Matthias <xmatthias@outlook.com>
Date: Thu, 24 Oct 2024 19:22:46 +0200
Subject: [PATCH] feat: add ProgressLike protocol

---
 freqtrade/util/__init__.py         |  6 +++++-
 freqtrade/util/progress_tracker.py | 18 +++++++++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/freqtrade/util/__init__.py b/freqtrade/util/__init__.py
index f3674c194..3d227880e 100644
--- a/freqtrade/util/__init__.py
+++ b/freqtrade/util/__init__.py
@@ -15,7 +15,11 @@ from freqtrade.util.formatters import decimals_per_coin, fmt_coin, fmt_coin2, ro
 from freqtrade.util.ft_precise import FtPrecise
 from freqtrade.util.measure_time import MeasureTime
 from freqtrade.util.periodic_cache import PeriodicCache
-from freqtrade.util.progress_tracker import get_progress_tracker  # noqa F401
+from freqtrade.util.progress_tracker import (  # noqa F401
+    ProgressLike,
+    get_progress_tracker,
+    retrieve_progress_tracker,
+)
 from freqtrade.util.rich_progress import CustomProgress
 from freqtrade.util.rich_tables import print_df_rich_table, print_rich_table
 from freqtrade.util.template_renderer import render_template, render_template_with_fallback  # noqa
diff --git a/freqtrade/util/progress_tracker.py b/freqtrade/util/progress_tracker.py
index ca8f5d3e9..2e339c7fb 100644
--- a/freqtrade/util/progress_tracker.py
+++ b/freqtrade/util/progress_tracker.py
@@ -1,6 +1,10 @@
+from contextlib import AbstractContextManager
+from typing import Optional, Protocol
+
 from rich.progress import (
     BarColumn,
     MofNCompleteColumn,
+    TaskID,
     TaskProgressColumn,
     TextColumn,
     TimeElapsedColumn,
@@ -10,7 +14,19 @@ from rich.progress import (
 from freqtrade.util.rich_progress import CustomProgress
 
 
-def get_progress_tracker(**kwargs):
+class ProgressLike(Protocol, AbstractContextManager["ProgressLike"]):
+    def add_task(self, description: str, *args, **kwargs) -> TaskID: ...
+
+    def update(self, task_id: TaskID, *, advance: float | None = None, **kwargs): ...
+
+
+def retrieve_progress_tracker(pt: Optional[ProgressLike]) -> ProgressLike:
+    if pt is None:
+        return get_progress_tracker()
+    return pt
+
+
+def get_progress_tracker(**kwargs) -> ProgressLike:
     """
     Get progress Bar with custom columns.
     """