From 360fe578d762625927ef1f2213c071b4d127224f Mon Sep 17 00:00:00 2001
From: Daniil Fajnberg <mail@daniil.fajnberg.de>
Date: Mon, 7 Mar 2022 14:34:40 +0100
Subject: [PATCH] additional base commands for control server

---
 src/asyncio_taskpool/constants.py | 11 +++++++++++
 src/asyncio_taskpool/session.py   | 27 ++++++++++++++++++++++-----
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/asyncio_taskpool/constants.py b/src/asyncio_taskpool/constants.py
index 9d2b749..5266469 100644
--- a/src/asyncio_taskpool/constants.py
+++ b/src/asyncio_taskpool/constants.py
@@ -37,10 +37,21 @@ class CLIENT_INFO:
 
 class CMD:
     __slots__ = ()
+    # Base commands:
     CMD = 'command'
     NAME = 'name'
     POOL_SIZE = 'pool-size'
+    IS_LOCKED = 'is-locked'
+    LOCK = 'lock'
+    UNLOCK = 'unlock'
     NUM_RUNNING = 'num-running'
+    NUM_CANCELLATIONS = 'num-cancellations'
+    NUM_ENDED = 'num-ended'
+    NUM_FINISHED = 'num-finished'
+    IS_FULL = 'is-full'
+    GET_GROUP_IDS = 'get-group-ids'
+
+    # Simple commands:
     START = 'start'
     STOP = 'stop'
     STOP_ALL = 'stop-all'
diff --git a/src/asyncio_taskpool/session.py b/src/asyncio_taskpool/session.py
index 2216076..71f6486 100644
--- a/src/asyncio_taskpool/session.py
+++ b/src/asyncio_taskpool/session.py
@@ -23,7 +23,7 @@ import logging
 import json
 from argparse import ArgumentError, HelpFormatter
 from asyncio.streams import StreamReader, StreamWriter
-from typing import Callable, Optional, Union, TYPE_CHECKING
+from typing import Callable, Optional, Type, Union, TYPE_CHECKING
 
 from .constants import CMD, SESSION_WRITER, SESSION_MSG_BYTES, CLIENT_INFO
 from .exceptions import HelpRequested, NotATaskPool, UnknownTaskPoolClass
@@ -108,19 +108,36 @@ class ControlSession:
         These include commands mapping to the following pool methods:
             - __str__
             - pool_size (get/set property)
+            - is_locked
+            - lock & unlock
             - num_running
         """
-        self._add_command(CMD.NAME, short_help=get_first_doc_line(self._pool.__class__.__str__))
+        cls: Type[BaseTaskPool] = self._pool.__class__
+        self._add_command(CMD.NAME, short_help=get_first_doc_line(cls.__str__))
         self._add_command(
             CMD.POOL_SIZE, 
             short_help="Get/set the maximum number of tasks in the pool.", 
             formatter_class=HelpFormatter
         ).add_optional_num_argument(
             default=None,
-            help=f"If passed a number: {get_first_doc_line(self._pool.__class__.pool_size.fset)} "
-                 f"If omitted: {get_first_doc_line(self._pool.__class__.pool_size.fget)}"
+            help=f"If passed a number: {get_first_doc_line(cls.pool_size.fset)} "
+                 f"If omitted: {get_first_doc_line(cls.pool_size.fget)}"
+        )
+        self._add_command(CMD.IS_LOCKED, short_help=get_first_doc_line(cls.is_locked.fget))
+        self._add_command(CMD.LOCK, short_help=get_first_doc_line(cls.lock))
+        self._add_command(CMD.UNLOCK, short_help=get_first_doc_line(cls.unlock))
+        self._add_command(CMD.NUM_RUNNING, short_help=get_first_doc_line(cls.num_running.fget))
+        self._add_command(CMD.NUM_CANCELLATIONS, short_help=get_first_doc_line(cls.num_cancellations.fget))
+        self._add_command(CMD.NUM_ENDED, short_help=get_first_doc_line(cls.num_ended.fget))
+        self._add_command(CMD.NUM_FINISHED, short_help=get_first_doc_line(cls.num_finished.fget))
+        self._add_command(CMD.IS_FULL, short_help=get_first_doc_line(cls.is_full.fget))
+        self._add_command(
+            CMD.GET_GROUP_IDS, short_help=get_first_doc_line(cls.get_group_ids)
+        ).add_argument(
+            'group_name',
+            nargs='*',
+            help="Must be a name of a task group that exists within the pool."
         )
-        self._add_command(CMD.NUM_RUNNING, short_help=get_first_doc_line(self._pool.__class__.num_running.fget))
 
     def _add_simple_commands(self) -> None:
         """