From 7e127cd5cce37b6c0f6cd8b0139ae9d13cb69e07 Mon Sep 17 00:00:00 2001
From: Joscha <joscha@plugh.de>
Date: Thu, 29 Apr 2021 13:43:50 +0200
Subject: [PATCH] Clean up and fix conductor and limiter

Turns out you have to await an async lock, who knew...
---
 PFERD/conductor.py | 33 ++++++++++++---------------------
 PFERD/limiter.py   |  8 ++------
 2 files changed, 14 insertions(+), 27 deletions(-)

diff --git a/PFERD/conductor.py b/PFERD/conductor.py
index 86df7e4..fef5a0e 100644
--- a/PFERD/conductor.py
+++ b/PFERD/conductor.py
@@ -1,12 +1,8 @@
 import asyncio
 from contextlib import asynccontextmanager, contextmanager
-from pathlib import Path
-# TODO If we upgrade to python 3.9, these context manager hints are deprecated
-from typing import (AsyncContextManager, AsyncIterator, ContextManager,
-                    Iterator, List, Optional)
+from typing import AsyncIterator, Iterator, List, Optional
 
 import rich
-from rich.markup import escape
 from rich.progress import Progress, TaskID
 
 
@@ -38,11 +34,11 @@ class TerminalConductor:
         self._stopped = True
 
     async def start(self) -> None:
-        with self._lock:
+        async with self._lock:
             self._start()
 
     async def stop(self) -> None:
-        with self._lock:
+        async with self._lock:
             self._stop()
 
     def print(self, line: str) -> None:
@@ -52,7 +48,7 @@ class TerminalConductor:
             rich.print(line)
 
     @asynccontextmanager
-    async def _exclusive_output_cm(self) -> AsyncIterator[None]:
+    async def exclusive_output(self) -> AsyncIterator[None]:
         async with self._lock:
             self.stop()
             try:
@@ -60,25 +56,20 @@ class TerminalConductor:
             finally:
                 self.start()
 
-    def exclusive_output(self) -> AsyncContextManager[None]:
-        return self._exclusive_output_cm()
-
     @contextmanager
-    def _progress_bar_cm(
+    def progress_bar(
             self,
             description: str,
-            steps: Optional[float],
+            total: Optional[float] = None,
     ) -> Iterator[ProgressBar]:
-        taskid = self._progress.add_task(description, steps=steps)
+        if total is None:
+            # Indeterminate progress bar
+            taskid = self._progress.add_task(description, start=False)
+        else:
+            taskid = self._progress.add_task(description, total=total)
+
         bar = ProgressBar(self._progress, taskid)
         try:
             yield bar
         finally:
             self._progress.remove_task(taskid)
-
-    def progress_bar(
-            self,
-            description: Path,
-            steps: Optional[float],
-    ) -> ContextManager[ProgressBar]:
-        return self._progress_bar_cm(escape(str(description)), steps=steps)
diff --git a/PFERD/limiter.py b/PFERD/limiter.py
index f001d8b..ff91d57 100644
--- a/PFERD/limiter.py
+++ b/PFERD/limiter.py
@@ -1,7 +1,6 @@
 import asyncio
 from contextlib import asynccontextmanager
-# TODO If we upgrade to python 3.9, this context manager hint is deprecated
-from typing import AsyncContextManager, AsyncIterator
+from typing import AsyncIterator
 
 
 class Limiter:
@@ -9,12 +8,9 @@ class Limiter:
         self._semaphore = asyncio.Semaphore(limit)
 
     @asynccontextmanager
-    async def _context_manager(self) -> AsyncIterator[None]:
+    async def limit(self) -> AsyncIterator[None]:
         await self._semaphore.acquire()
         try:
             yield
         finally:
             self._semaphore.release()
-
-    def limit(self) -> AsyncContextManager[None]:
-        return self._context_manager()