generated from daniil-berg/boilerplate-py
big rework of the session-parser-interaction;
dynamically adding pool methods/properties as parser commands; dynamically executing selected pool method/property; greatly simplified `ControlSession` class; removed the need for hard-coded command names; adjusted unittests accordingly
This commit is contained in:
@ -27,32 +27,12 @@ DATETIME_FORMAT = '%Y-%m-%d_%H-%M-%S'
|
||||
CLIENT_EXIT = 'exit'
|
||||
|
||||
SESSION_MSG_BYTES = 1024 * 100
|
||||
SESSION_WRITER = 'session_writer'
|
||||
|
||||
STREAM_WRITER = 'stream_writer'
|
||||
CMD = 'command'
|
||||
CMD_OK = b"ok"
|
||||
|
||||
|
||||
class CLIENT_INFO:
|
||||
__slots__ = ()
|
||||
TERMINAL_WIDTH = 'terminal_width'
|
||||
|
||||
|
||||
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'
|
||||
FUNC_NAME = 'func-name'
|
||||
|
||||
@ -63,13 +63,5 @@ class ServerException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class UnknownTaskPoolClass(ServerException):
|
||||
pass
|
||||
|
||||
|
||||
class NotATaskPool(ServerException):
|
||||
pass
|
||||
|
||||
|
||||
class HelpRequested(ServerException):
|
||||
pass
|
||||
|
||||
@ -51,10 +51,6 @@ async def join_queue(q: Queue) -> None:
|
||||
await q.join()
|
||||
|
||||
|
||||
def tasks_str(num: int) -> str:
|
||||
return "tasks" if num != 1 else "task"
|
||||
|
||||
|
||||
def get_first_doc_line(obj: object) -> str:
|
||||
return getdoc(obj).strip().split("\n", 1)[0].strip()
|
||||
|
||||
|
||||
299
src/asyncio_taskpool/parser.py
Normal file
299
src/asyncio_taskpool/parser.py
Normal file
@ -0,0 +1,299 @@
|
||||
__author__ = "Daniil Fajnberg"
|
||||
__copyright__ = "Copyright © 2022 Daniil Fajnberg"
|
||||
__license__ = """GNU LGPLv3.0
|
||||
|
||||
This file is part of asyncio-taskpool.
|
||||
|
||||
asyncio-taskpool is free software: you can redistribute it and/or modify it under the terms of
|
||||
version 3.0 of the GNU Lesser General Public License as published by the Free Software Foundation.
|
||||
|
||||
asyncio-taskpool is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
|
||||
without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
|
||||
See the GNU Lesser General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU Lesser General Public License along with asyncio-taskpool.
|
||||
If not, see <https://www.gnu.org/licenses/>."""
|
||||
|
||||
__doc__ = """
|
||||
This module contains the the definition of the `ControlParser` class used by a control server.
|
||||
"""
|
||||
|
||||
|
||||
from argparse import Action, ArgumentParser, ArgumentDefaultsHelpFormatter, HelpFormatter, SUPPRESS
|
||||
from asyncio.streams import StreamWriter
|
||||
from inspect import Parameter, getmembers, isfunction, signature
|
||||
from shutil import get_terminal_size
|
||||
from typing import Callable, Container, Dict, Set, Type, TypeVar
|
||||
|
||||
from .constants import CLIENT_INFO, CMD, STREAM_WRITER
|
||||
from .exceptions import HelpRequested
|
||||
from .helpers import get_first_doc_line
|
||||
|
||||
|
||||
FmtCls = TypeVar('FmtCls', bound=Type[HelpFormatter])
|
||||
ParsersDict = Dict[str, 'ControlParser']
|
||||
|
||||
OMIT_PARAMS_DEFAULT = ('self', )
|
||||
|
||||
FORMATTER_CLASS = 'formatter_class'
|
||||
NAME, PROG, HELP, DESCRIPTION = 'name', 'prog', 'help', 'description'
|
||||
|
||||
|
||||
class ControlParser(ArgumentParser):
|
||||
"""
|
||||
Subclass of the standard `argparse.ArgumentParser` for remote interaction.
|
||||
|
||||
Such a parser is not supposed to ever print to stdout/stderr, but instead direct all messages to a `StreamWriter`
|
||||
instance passed to it during initialization.
|
||||
Furthermore, it requires defining the width of the terminal, to adjust help formatting to the terminal size of a
|
||||
connected client.
|
||||
Finally, it offers some convenience methods and makes use of custom exceptions.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def help_formatter_factory(terminal_width: int, base_cls: FmtCls = None) -> FmtCls:
|
||||
"""
|
||||
Constructs and returns a subclass of `argparse.HelpFormatter` with a fixed terminal width argument.
|
||||
|
||||
Although a custom formatter class can be explicitly passed into the `ArgumentParser` constructor, this is not
|
||||
as convenient, when making use of sub-parsers.
|
||||
|
||||
Args:
|
||||
terminal_width:
|
||||
The number of columns of the terminal to which to adjust help formatting.
|
||||
base_cls (optional):
|
||||
The base class to use for inheritance. By default `argparse.ArgumentDefaultsHelpFormatter` is used.
|
||||
|
||||
Returns:
|
||||
The subclass of `base_cls` which fixes the constructor's `width` keyword-argument to `terminal_width`.
|
||||
"""
|
||||
if base_cls is None:
|
||||
base_cls = ArgumentDefaultsHelpFormatter
|
||||
|
||||
class ClientHelpFormatter(base_cls):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
kwargs['width'] = terminal_width
|
||||
super().__init__(*args, **kwargs)
|
||||
return ClientHelpFormatter
|
||||
|
||||
def __init__(self, stream_writer: StreamWriter, terminal_width: int = None,
|
||||
**kwargs) -> None:
|
||||
"""
|
||||
Sets additional internal attributes depending on whether a parent-parser was defined.
|
||||
|
||||
The `help_formatter_factory` is called and the returned class is mapped to the `FORMATTER_CLASS` keyword.
|
||||
By default, `exit_on_error` is set to `False` (as opposed to how the parent class handles it).
|
||||
|
||||
Args:
|
||||
stream_writer:
|
||||
The instance of the `asyncio.StreamWriter` to use for message output.
|
||||
terminal_width (optional):
|
||||
The terminal width to assume for all message formatting. Defaults to `shutil.get_terminal_size`.
|
||||
**kwargs(optional):
|
||||
In addition to the regular `ArgumentParser` constructor parameters, this method expects the instance of
|
||||
the `StreamWriter` as well as the terminal width both to be passed explicitly, if the `parent` argument
|
||||
is empty.
|
||||
"""
|
||||
self._stream_writer: StreamWriter = stream_writer
|
||||
self._terminal_width: int = terminal_width if terminal_width is not None else get_terminal_size().columns
|
||||
kwargs[FORMATTER_CLASS] = self.help_formatter_factory(self._terminal_width, kwargs.get(FORMATTER_CLASS))
|
||||
kwargs.setdefault('exit_on_error', False)
|
||||
super().__init__(**kwargs)
|
||||
self._flags: Set[str] = set()
|
||||
self._commands = None
|
||||
|
||||
def add_function_command(self, function: Callable, omit_params: Container[str] = OMIT_PARAMS_DEFAULT,
|
||||
**subparser_kwargs) -> 'ControlParser':
|
||||
"""
|
||||
Takes a function along with its parameters and adds a corresponding (sub-)command to the parser.
|
||||
|
||||
The `add_subparsers` method must have been called prior to this.
|
||||
|
||||
NOTE: Currently, only a limited spectrum of parameters can be accurately converted to a parser argument.
|
||||
This method works correctly with any public method of the `SimpleTaskPool` class.
|
||||
|
||||
Args:
|
||||
function:
|
||||
The reference to the function to be "converted" to a parser command.
|
||||
omit_params (optional):
|
||||
Names of function parameters not to add as parser arguments.
|
||||
**subparser_kwargs (optional):
|
||||
Passed directly to the `add_parser` method.
|
||||
|
||||
Returns:
|
||||
The subparser instance created from the function.
|
||||
"""
|
||||
subparser_kwargs.setdefault(NAME, function.__name__.replace('_', '-'))
|
||||
subparser_kwargs.setdefault(PROG, subparser_kwargs[NAME])
|
||||
subparser_kwargs.setdefault(HELP, get_first_doc_line(function))
|
||||
subparser_kwargs.setdefault(DESCRIPTION, subparser_kwargs[HELP])
|
||||
subparser: ControlParser = self._commands.add_parser(**subparser_kwargs)
|
||||
subparser.add_function_args(function, omit_params)
|
||||
return subparser
|
||||
|
||||
def add_property_command(self, prop: property, cls_name: str = '', **subparser_kwargs) -> 'ControlParser':
|
||||
"""
|
||||
Same as the `add_function_command` method, but for properties.
|
||||
|
||||
Args:
|
||||
prop:
|
||||
The reference to the property to be "converted" to a parser command.
|
||||
cls_name (optional):
|
||||
Name of the class the property is defined on to appear in the command help text.
|
||||
**subparser_kwargs (optional):
|
||||
Passed directly to the `add_parser` method.
|
||||
|
||||
Returns:
|
||||
The subparser instance created from the property.
|
||||
"""
|
||||
subparser_kwargs.setdefault(NAME, prop.fget.__name__.replace('_', '-'))
|
||||
subparser_kwargs.setdefault(PROG, subparser_kwargs[NAME])
|
||||
getter_help = get_first_doc_line(prop.fget)
|
||||
if prop.fset is None:
|
||||
subparser_kwargs.setdefault(HELP, getter_help)
|
||||
else:
|
||||
subparser_kwargs.setdefault(HELP, f"Get/set the `{cls_name}.{subparser_kwargs[NAME]}` property")
|
||||
subparser_kwargs.setdefault(DESCRIPTION, subparser_kwargs[HELP])
|
||||
subparser: ControlParser = self._commands.add_parser(**subparser_kwargs)
|
||||
if prop.fset is not None:
|
||||
_, param = signature(prop.fset).parameters.values()
|
||||
setter_arg_help = f"If provided: {get_first_doc_line(prop.fset)} If omitted: {getter_help}"
|
||||
subparser.add_function_arg(param, nargs='?', default=SUPPRESS, help=setter_arg_help)
|
||||
return subparser
|
||||
|
||||
def add_class_commands(self, cls: Type, public_only: bool = True, omit_members: Container[str] = (),
|
||||
member_arg_name: str = CMD) -> ParsersDict:
|
||||
"""
|
||||
Takes a class and adds its methods and properties as (sub-)commands to the parser.
|
||||
|
||||
The `add_subparsers` method must have been called prior to this.
|
||||
|
||||
NOTE: Currently, only a limited spectrum of function parameters can be accurately converted to parser arguments.
|
||||
This method works correctly with the `SimpleTaskPool` class.
|
||||
|
||||
Args:
|
||||
cls:
|
||||
The reference to the class whose methods/properties are to be "converted" to parser commands.
|
||||
public_only (optional):
|
||||
If `False`, protected and private members are considered as well. `True` by default.
|
||||
omit_members (optional):
|
||||
Names of functions/properties not to add as parser commands.
|
||||
member_arg_name (optional):
|
||||
After parsing the arguments, depending on which command was invoked by the user, the corresponding
|
||||
method/property will be stored as an extra argument in the parsed namespace under this attribute name.
|
||||
Defaults to `constants.CMD`.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping class member names to the (sub-)parsers created from them.
|
||||
"""
|
||||
parsers: ParsersDict = {}
|
||||
common_kwargs = {STREAM_WRITER: self._stream_writer, CLIENT_INFO.TERMINAL_WIDTH: self._terminal_width}
|
||||
for name, member in getmembers(cls):
|
||||
if name in omit_members or (name.startswith('_') and public_only):
|
||||
continue
|
||||
if isfunction(member):
|
||||
subparser = self.add_function_command(member, **common_kwargs)
|
||||
elif isinstance(member, property):
|
||||
subparser = self.add_property_command(member, cls.__name__, **common_kwargs)
|
||||
else:
|
||||
continue
|
||||
subparser.set_defaults(**{member_arg_name: member})
|
||||
parsers[name] = subparser
|
||||
return parsers
|
||||
|
||||
def add_subparsers(self, *args, **kwargs):
|
||||
"""Adds the subparsers action as an internal attribute before returning it."""
|
||||
self._commands = super().add_subparsers(*args, **kwargs)
|
||||
return self._commands
|
||||
|
||||
def _print_message(self, message: str, *args, **kwargs) -> None:
|
||||
"""This is overridden to ensure that no messages are sent to stdout/stderr, but always to the stream writer."""
|
||||
if message:
|
||||
self._stream_writer.write(message.encode())
|
||||
|
||||
def exit(self, status: int = 0, message: str = None) -> None:
|
||||
"""This is overridden to prevent system exit to be invoked."""
|
||||
if message:
|
||||
self._print_message(message)
|
||||
|
||||
def error(self, message: str) -> None:
|
||||
"""This just adds the custom `HelpRequested` exception after the parent class' method."""
|
||||
super().error(message=message)
|
||||
raise HelpRequested
|
||||
|
||||
def print_help(self, file=None) -> None:
|
||||
"""This just adds the custom `HelpRequested` exception after the parent class' method."""
|
||||
super().print_help(file)
|
||||
raise HelpRequested
|
||||
|
||||
def add_function_arg(self, parameter: Parameter, **kwargs) -> Action:
|
||||
"""
|
||||
Takes an `inspect.Parameter` of a function and adds a corresponding argument to the parser.
|
||||
|
||||
NOTE: Currently, only a limited spectrum of parameters can be accurately converted to a parser argument.
|
||||
This method works correctly with any parameter of any public method of the `SimpleTaskPool` class.
|
||||
|
||||
Args:
|
||||
parameter: The `inspect.Parameter` object to be converted to a parser argument.
|
||||
**kwargs: Passed to the `add_argument` method of the base class.
|
||||
|
||||
Returns:
|
||||
The `argparse.Action` returned by the `add_argument` method.
|
||||
"""
|
||||
if parameter.default is Parameter.empty:
|
||||
# A non-optional function parameter should correspond to a positional argument.
|
||||
name_or_flags = [parameter.name]
|
||||
else:
|
||||
flag = None
|
||||
long = f'--{parameter.name.replace("_", "-")}'
|
||||
# We try to generate a short version (flag) for the argument.
|
||||
letter = parameter.name[0]
|
||||
if letter not in self._flags:
|
||||
flag = f'-{letter}'
|
||||
self._flags.add(letter)
|
||||
elif letter.upper() not in self._flags:
|
||||
flag = f'-{letter.upper()}'
|
||||
self._flags.add(letter.upper())
|
||||
name_or_flags = [long] if flag is None else [flag, long]
|
||||
if parameter.annotation is bool:
|
||||
# If we are dealing with a boolean parameter, always use the 'store_true' action.
|
||||
# Even if the parameter's default value is `True`, this will make the parser argument's default `False`.
|
||||
kwargs.setdefault('action', 'store_true')
|
||||
else:
|
||||
# For now, any other type annotation will implicitly use the default action 'store'.
|
||||
# In addition, we always set the default value.
|
||||
kwargs.setdefault('default', parameter.default)
|
||||
if parameter.kind == Parameter.VAR_POSITIONAL:
|
||||
# This is to be able to later unpack an arbitrary number of positional arguments.
|
||||
kwargs.setdefault('nargs', '*')
|
||||
if not kwargs.get('action') == 'store_true':
|
||||
# The lambda wrapper around the type annotation is to avoid ValueError being raised on suppressed arguments.
|
||||
# See: https://bugs.python.org/issue36078
|
||||
kwargs.setdefault('type', get_arg_type_wrapper(parameter.annotation))
|
||||
return self.add_argument(*name_or_flags, **kwargs)
|
||||
|
||||
def add_function_args(self, function: Callable, omit: Container[str] = OMIT_PARAMS_DEFAULT) -> None:
|
||||
"""
|
||||
Takes a function reference and adds its parameters as arguments to the parser.
|
||||
|
||||
NOTE: Currently, only a limited spectrum of parameters can be accurately converted to a parser argument.
|
||||
This method works correctly with any public method of the `SimpleTaskPool` class.
|
||||
|
||||
Args:
|
||||
function:
|
||||
The function whose parameters are to be converted to parser arguments.
|
||||
Its parameters must be properly annotated.
|
||||
omit (optional):
|
||||
Names of function parameters not to add as parser arguments.
|
||||
"""
|
||||
for param in signature(function).parameters.values():
|
||||
if param.name not in omit:
|
||||
# TODO: Look into parsing docstrings properly to try and extract argument help text.
|
||||
# For now, the argument help just shows the type it will be converted to.
|
||||
self.add_function_arg(param, help=repr(param.annotation))
|
||||
|
||||
|
||||
def get_arg_type_wrapper(cls: Type) -> Callable:
|
||||
def wrapper(arg):
|
||||
return arg if arg is SUPPRESS else cls(arg)
|
||||
return wrapper
|
||||
@ -15,21 +15,22 @@ You should have received a copy of the GNU Lesser General Public License along w
|
||||
If not, see <https://www.gnu.org/licenses/>."""
|
||||
|
||||
__doc__ = """
|
||||
This module contains the the definition of the control session class used by the control server.
|
||||
This module contains the the definition of the `ControlSession` class used by the control server.
|
||||
"""
|
||||
|
||||
|
||||
import logging
|
||||
import json
|
||||
from argparse import ArgumentError, HelpFormatter
|
||||
from argparse import ArgumentError
|
||||
from asyncio.streams import StreamReader, StreamWriter
|
||||
from typing import Callable, Optional, Type, Union, TYPE_CHECKING
|
||||
from inspect import isfunction, signature
|
||||
from typing import Callable, Optional, Union, TYPE_CHECKING
|
||||
|
||||
from .constants import CMD, SESSION_WRITER, SESSION_MSG_BYTES, CLIENT_INFO
|
||||
from .exceptions import HelpRequested, NotATaskPool, UnknownTaskPoolClass
|
||||
from .helpers import get_first_doc_line, return_or_exception, tasks_str
|
||||
from .pool import BaseTaskPool, TaskPool, SimpleTaskPool
|
||||
from .session_parser import CommandParser, NUM
|
||||
from .constants import CLIENT_INFO, CMD, CMD_OK, SESSION_MSG_BYTES, STREAM_WRITER
|
||||
from .exceptions import HelpRequested
|
||||
from .helpers import return_or_exception
|
||||
from .pool import TaskPool, SimpleTaskPool
|
||||
from .parser import ControlParser
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .server import ControlServer
|
||||
@ -67,273 +68,88 @@ class ControlSession:
|
||||
self._client_class_name = server.client_class_name
|
||||
self._reader: StreamReader = reader
|
||||
self._writer: StreamWriter = writer
|
||||
self._parser: Optional[CommandParser] = None
|
||||
self._subparsers = None
|
||||
self._parser: Optional[ControlParser] = None
|
||||
|
||||
def _add_command(self, name: str, prog: str = None, short_help: str = None, long_help: str = None,
|
||||
**kwargs) -> CommandParser:
|
||||
async def _exec_method_and_respond(self, method: Callable, **kwargs) -> None:
|
||||
"""
|
||||
Convenience method for adding a subparser (i.e. another command) to the main `CommandParser` instance.
|
||||
Takes a pool method reference, executes it, and writes a response accordingly.
|
||||
|
||||
Will always pass the session's main `CommandParser` instance as the `parent` keyword-argument.
|
||||
If the first parameter is named `self`, the method will be called with the `_pool` instance as its first
|
||||
positional argument. If it returns nothing, the response upon successful execution will be `constants.CMD_OK`,
|
||||
otherwise the response written to the stream will be its return value (as an encoded string).
|
||||
|
||||
Args:
|
||||
name:
|
||||
The command name; passed directly into the `add_parser` method.
|
||||
prog (optional):
|
||||
Also passed into the `add_parser` method as the corresponding keyword-argument. By default, is set
|
||||
equal to the `name` argument.
|
||||
short_help (optional):
|
||||
Passed into the `add_parser` method as the `help` keyword-argument, unless it is left empty and the
|
||||
`long_help` argument is present; in that case the `long_help` argument is passed as `help`.
|
||||
long_help (optional):
|
||||
Passed into the `add_parser` method as the `description` keyword-argument, unless it is left empty and
|
||||
the `short_help` argument is present; in that case the `short_help` argument is passed as `description`.
|
||||
prop:
|
||||
The reference to the method defined on the `_pool` instance's class.
|
||||
**kwargs (optional):
|
||||
Any keyword-arguments to directly pass into the `add_parser` method.
|
||||
|
||||
Returns:
|
||||
An instance of the `CommandParser` class representing the newly added control command.
|
||||
Must correspond to the arguments expected by the `method`.
|
||||
Correctly unpacks arbitrary-length positional and keyword-arguments.
|
||||
"""
|
||||
if prog is None:
|
||||
prog = name
|
||||
kwargs.setdefault('help', short_help or long_help)
|
||||
kwargs.setdefault('description', long_help or short_help)
|
||||
return self._subparsers.add_parser(name, prog=prog, parent=self._parser, **kwargs)
|
||||
log.warning("%s calls %s.%s", self._client_class_name, self._pool.__class__.__name__, method.__name__)
|
||||
normal_pos, var_pos = [], []
|
||||
for param in signature(method).parameters.values():
|
||||
if param.name == 'self':
|
||||
normal_pos.append(self._pool)
|
||||
elif param.kind in (param.POSITIONAL_OR_KEYWORD, param.POSITIONAL_ONLY):
|
||||
normal_pos.append(kwargs.pop(param.name))
|
||||
elif param.kind == param.VAR_POSITIONAL:
|
||||
var_pos = kwargs.pop(param.name)
|
||||
output = await return_or_exception(method, *normal_pos, *var_pos, **kwargs)
|
||||
self._writer.write(CMD_OK if output is None else str(output).encode())
|
||||
|
||||
def _add_base_commands(self) -> None:
|
||||
async def _exec_property_and_respond(self, prop: property, **kwargs) -> None:
|
||||
"""
|
||||
Adds the commands that are supported regardless of the specific subclass of `BaseTaskPool` controlled.
|
||||
Takes a pool property reference, executes its setter or getter, and writes a response accordingly.
|
||||
|
||||
These include commands mapping to the following pool methods:
|
||||
- __str__
|
||||
- pool_size (get/set property)
|
||||
- is_locked
|
||||
- lock & unlock
|
||||
- num_running
|
||||
"""
|
||||
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(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."
|
||||
)
|
||||
|
||||
def _add_simple_commands(self) -> None:
|
||||
"""
|
||||
Adds the commands that are only supported, if a `SimpleTaskPool` object is controlled.
|
||||
|
||||
These include commands mapping to the following pool methods:
|
||||
- start
|
||||
- stop
|
||||
- stop_all
|
||||
- func_name
|
||||
"""
|
||||
self._add_command(
|
||||
CMD.START, short_help=get_first_doc_line(self._pool.__class__.start)
|
||||
).add_optional_num_argument(
|
||||
help="Number of tasks to start."
|
||||
)
|
||||
self._add_command(
|
||||
CMD.STOP, short_help=get_first_doc_line(self._pool.__class__.stop)
|
||||
).add_optional_num_argument(
|
||||
help="Number of tasks to stop."
|
||||
)
|
||||
self._add_command(CMD.STOP_ALL, short_help=get_first_doc_line(self._pool.__class__.stop_all))
|
||||
self._add_command(CMD.FUNC_NAME, short_help=get_first_doc_line(self._pool.__class__.func_name.fget))
|
||||
|
||||
def _add_advanced_commands(self) -> None:
|
||||
"""
|
||||
Adds the commands that are only supported, if a `TaskPool` object is controlled.
|
||||
|
||||
These include commands mapping to the following pool methods:
|
||||
- ...
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def _init_parser(self, client_terminal_width: int) -> None:
|
||||
"""
|
||||
Initializes and fully configures the `CommandParser` responsible for handling the input.
|
||||
|
||||
Depending on what specific task pool class is controlled by the server, different commands are added.
|
||||
The property set/get method will always be called with the `_pool` instance as its first positional argument.
|
||||
|
||||
Args:
|
||||
client_terminal_width:
|
||||
The number of columns of the client's terminal to be able to nicely format messages from the parser.
|
||||
prop:
|
||||
The reference to the property defined on the `_pool` instance's class.
|
||||
**kwargs (optional):
|
||||
If not empty, the property setter is executed and the keyword arguments are passed along to it; the
|
||||
response upon successful execution will be `constants.CMD_OK`. Otherwise the property getter is
|
||||
executed and the response written to the stream will be its return value (as an encoded string).
|
||||
"""
|
||||
parser_kwargs = {
|
||||
'prog': '',
|
||||
SESSION_WRITER: self._writer,
|
||||
CLIENT_INFO.TERMINAL_WIDTH: client_terminal_width,
|
||||
}
|
||||
self._parser = CommandParser(**parser_kwargs)
|
||||
self._subparsers = self._parser.add_subparsers(title="Commands", dest=CMD.CMD)
|
||||
self._add_base_commands()
|
||||
if isinstance(self._pool, TaskPool):
|
||||
self._add_advanced_commands()
|
||||
elif isinstance(self._pool, SimpleTaskPool):
|
||||
self._add_simple_commands()
|
||||
elif isinstance(self._pool, BaseTaskPool):
|
||||
raise UnknownTaskPoolClass(f"No interface defined for {self._pool.__class__.__name__}")
|
||||
if kwargs:
|
||||
log.warning("%s sets %s.%s", self._client_class_name, self._pool.__class__.__name__, prop.fset.__name__)
|
||||
await return_or_exception(prop.fset, self._pool, **kwargs)
|
||||
self._writer.write(CMD_OK)
|
||||
else:
|
||||
raise NotATaskPool(f"Not a task pool instance: {self._pool}")
|
||||
log.warning("%s gets %s.%s", self._client_class_name, self._pool.__class__.__name__, prop.fget.__name__)
|
||||
self._writer.write(str(await return_or_exception(prop.fget, self._pool)).encode())
|
||||
|
||||
async def client_handshake(self) -> None:
|
||||
"""
|
||||
This method must be invoked before starting any other client interaction.
|
||||
|
||||
Client info is retrieved, server info is sent back, and the `CommandParser` is initialized and configured.
|
||||
Client info is retrieved, server info is sent back, and the `ControlParser` is initialized and configured.
|
||||
"""
|
||||
client_info = json.loads((await self._reader.read(SESSION_MSG_BYTES)).decode().strip())
|
||||
log.debug("%s connected", self._client_class_name)
|
||||
self._init_parser(client_info[CLIENT_INFO.TERMINAL_WIDTH])
|
||||
parser_kwargs = {
|
||||
STREAM_WRITER: self._writer,
|
||||
CLIENT_INFO.TERMINAL_WIDTH: client_info[CLIENT_INFO.TERMINAL_WIDTH],
|
||||
'prog': '',
|
||||
'usage': f'%(prog)s [-h] [{CMD}] ...'
|
||||
}
|
||||
self._parser = ControlParser(**parser_kwargs)
|
||||
self._parser.add_subparsers(title="Commands",
|
||||
metavar="(A command followed by '-h' or '--help' will show command-specific help.)")
|
||||
self._parser.add_class_commands(self._pool.__class__)
|
||||
self._writer.write(str(self._pool).encode())
|
||||
await self._writer.drain()
|
||||
|
||||
async def _write_function_output(self, func: Callable, *args, **kwargs) -> None:
|
||||
"""
|
||||
Acts as a wrapper around a call to a specific task pool method.
|
||||
|
||||
The method is called and any exception is caught and saved. If there is no output and no exception caught, a
|
||||
generic confirmation message is sent back to the client. Otherwise the output or a string representation of
|
||||
the exception caught is sent back.
|
||||
|
||||
Args:
|
||||
func:
|
||||
Reference to the task pool method.
|
||||
*args (optional):
|
||||
Any positional arguments to call the method with.
|
||||
*+kwargs (optional):
|
||||
Any keyword-arguments to call the method with.
|
||||
"""
|
||||
output = await return_or_exception(func, *args, **kwargs)
|
||||
self._writer.write(b"ok" if output is None else str(output).encode())
|
||||
|
||||
async def _cmd_name(self, **_kwargs) -> None:
|
||||
"""Maps to the `__str__` method of any task pool class."""
|
||||
log.debug("%s requests task pool name", self._client_class_name)
|
||||
await self._write_function_output(self._pool.__class__.__str__, self._pool)
|
||||
|
||||
async def _cmd_pool_size(self, **kwargs) -> None:
|
||||
"""Maps to the `pool_size` property of any task pool class."""
|
||||
num = kwargs.get(NUM)
|
||||
if num is None:
|
||||
log.debug("%s requests pool size", self._client_class_name)
|
||||
await self._write_function_output(self._pool.__class__.pool_size.fget, self._pool)
|
||||
else:
|
||||
log.debug("%s requests setting pool size to %s", self._client_class_name, num)
|
||||
await self._write_function_output(self._pool.__class__.pool_size.fset, self._pool, num)
|
||||
|
||||
async def _cmd_is_locked(self, **_kwargs) -> None:
|
||||
"""Maps to the `is_locked` property of any task pool class."""
|
||||
log.debug("%s checks locked status", self._client_class_name)
|
||||
await self._write_function_output(self._pool.__class__.is_locked.fget, self._pool)
|
||||
|
||||
async def _cmd_lock(self, **_kwargs) -> None:
|
||||
"""Maps to the `lock` method of any task pool class."""
|
||||
log.debug("%s requests locking the pool", self._client_class_name)
|
||||
await self._write_function_output(self._pool.lock)
|
||||
|
||||
async def _cmd_unlock(self, **_kwargs) -> None:
|
||||
"""Maps to the `unlock` method of any task pool class."""
|
||||
log.debug("%s requests unlocking the pool", self._client_class_name)
|
||||
await self._write_function_output(self._pool.unlock)
|
||||
|
||||
async def _cmd_num_running(self, **_kwargs) -> None:
|
||||
"""Maps to the `num_running` property of any task pool class."""
|
||||
log.debug("%s requests number of running tasks", self._client_class_name)
|
||||
await self._write_function_output(self._pool.__class__.num_running.fget, self._pool)
|
||||
|
||||
async def _cmd_num_cancellations(self, **_kwargs) -> None:
|
||||
"""Maps to the `num_cancellations` property of any task pool class."""
|
||||
log.debug("%s requests number of cancelled tasks", self._client_class_name)
|
||||
await self._write_function_output(self._pool.__class__.num_cancellations.fget, self._pool)
|
||||
|
||||
async def _cmd_num_ended(self, **_kwargs) -> None:
|
||||
"""Maps to the `num_ended` property of any task pool class."""
|
||||
log.debug("%s requests number of ended tasks", self._client_class_name)
|
||||
await self._write_function_output(self._pool.__class__.num_ended.fget, self._pool)
|
||||
|
||||
async def _cmd_num_finished(self, **_kwargs) -> None:
|
||||
"""Maps to the `num_finished` property of any task pool class."""
|
||||
log.debug("%s requests number of finished tasks", self._client_class_name)
|
||||
await self._write_function_output(self._pool.__class__.num_finished.fget, self._pool)
|
||||
|
||||
async def _cmd_is_full(self, **_kwargs) -> None:
|
||||
"""Maps to the `is_full` property of any task pool class."""
|
||||
log.debug("%s checks full status", self._client_class_name)
|
||||
await self._write_function_output(self._pool.__class__.is_full.fget, self._pool)
|
||||
|
||||
async def _cmd_get_group_ids(self, **kwargs) -> None:
|
||||
"""Maps to the `get_group_ids` method of any task pool class."""
|
||||
log.debug("%s requests task ids for groups %s", self._client_class_name, kwargs['group_name'])
|
||||
await self._write_function_output(self._pool.get_group_ids, *kwargs['group_name'])
|
||||
|
||||
async def _cmd_start(self, **kwargs) -> None:
|
||||
"""Maps to the `start` method of the `SimpleTaskPool` class."""
|
||||
num = kwargs[NUM]
|
||||
log.debug("%s requests starting %s %s", self._client_class_name, num, tasks_str(num))
|
||||
await self._write_function_output(self._pool.start, num)
|
||||
|
||||
async def _cmd_stop(self, **kwargs) -> None:
|
||||
"""Maps to the `stop` method of the `SimpleTaskPool` class."""
|
||||
num = kwargs[NUM]
|
||||
log.debug("%s requests stopping %s %s", self._client_class_name, num, tasks_str(num))
|
||||
await self._write_function_output(self._pool.stop, num)
|
||||
|
||||
async def _cmd_stop_all(self, **_kwargs) -> None:
|
||||
"""Maps to the `stop_all` method of the `SimpleTaskPool` class."""
|
||||
log.debug("%s requests stopping all tasks", self._client_class_name)
|
||||
await self._write_function_output(self._pool.stop_all)
|
||||
|
||||
async def _cmd_func_name(self, **_kwargs) -> None:
|
||||
"""Maps to the `func_name` method of the `SimpleTaskPool` class."""
|
||||
log.debug("%s requests pool function name", self._client_class_name)
|
||||
await self._write_function_output(self._pool.__class__.func_name.fget, self._pool)
|
||||
|
||||
async def _execute_command(self, **kwargs) -> None:
|
||||
"""
|
||||
Dynamically gets the correct `_cmd_...` method depending on the name of the command passed and executes it.
|
||||
|
||||
Args:
|
||||
**kwargs:
|
||||
Must include the `CMD.CMD` key mapping the the command name. The rest of the keyword-arguments is
|
||||
simply passed into the method determined from the command name.
|
||||
"""
|
||||
method = getattr(self, f'_cmd_{kwargs.pop(CMD.CMD).replace("-", "_")}')
|
||||
await method(**kwargs)
|
||||
|
||||
async def _parse_command(self, msg: str) -> None:
|
||||
"""
|
||||
Takes a message from the client and attempts to parse it.
|
||||
|
||||
If a parsing error occurs, it is returned to the client. If the `HelpRequested` exception was raised by the
|
||||
`CommandParser`, nothing else happens. Otherwise, the `_execute_command` method is called with the entire
|
||||
dictionary of keyword-arguments returned by the `CommandParser` passed into it.
|
||||
`ControlParser`, nothing else happens. Otherwise, the appropriate `_exec...` method is called with the entire
|
||||
dictionary of keyword-arguments returned by the `ControlParser` passed into it.
|
||||
|
||||
Args:
|
||||
msg:
|
||||
The non-empty string read from the client stream.
|
||||
msg: The non-empty string read from the client stream.
|
||||
"""
|
||||
try:
|
||||
kwargs = vars(self._parser.parse_args(msg.split(' ')))
|
||||
@ -342,7 +158,11 @@ class ControlSession:
|
||||
return
|
||||
except HelpRequested:
|
||||
return
|
||||
await self._execute_command(**kwargs)
|
||||
command = kwargs.pop(CMD)
|
||||
if isfunction(command):
|
||||
await self._exec_method_and_respond(command, **kwargs)
|
||||
elif isinstance(command, property):
|
||||
await self._exec_property_and_respond(command, **kwargs)
|
||||
|
||||
async def listen(self) -> None:
|
||||
"""
|
||||
|
||||
@ -1,127 +0,0 @@
|
||||
__author__ = "Daniil Fajnberg"
|
||||
__copyright__ = "Copyright © 2022 Daniil Fajnberg"
|
||||
__license__ = """GNU LGPLv3.0
|
||||
|
||||
This file is part of asyncio-taskpool.
|
||||
|
||||
asyncio-taskpool is free software: you can redistribute it and/or modify it under the terms of
|
||||
version 3.0 of the GNU Lesser General Public License as published by the Free Software Foundation.
|
||||
|
||||
asyncio-taskpool is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
|
||||
without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
|
||||
See the GNU Lesser General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU Lesser General Public License along with asyncio-taskpool.
|
||||
If not, see <https://www.gnu.org/licenses/>."""
|
||||
|
||||
__doc__ = """
|
||||
This module contains the the definition of the `CommandParser` class used in a control server session.
|
||||
"""
|
||||
|
||||
|
||||
from argparse import Action, ArgumentParser, ArgumentDefaultsHelpFormatter, HelpFormatter
|
||||
from asyncio.streams import StreamWriter
|
||||
from typing import Type, TypeVar
|
||||
|
||||
from .constants import SESSION_WRITER, CLIENT_INFO
|
||||
from .exceptions import HelpRequested
|
||||
|
||||
|
||||
FmtCls = TypeVar('FmtCls', bound=Type[HelpFormatter])
|
||||
FORMATTER_CLASS = 'formatter_class'
|
||||
NUM = 'num'
|
||||
|
||||
|
||||
class CommandParser(ArgumentParser):
|
||||
"""
|
||||
Subclass of the standard `argparse.ArgumentParser` for remote interaction.
|
||||
|
||||
Such a parser is not supposed to ever print to stdout/stderr, but instead direct all messages to a `StreamWriter`
|
||||
instance passed to it during initialization.
|
||||
Furthermore, it requires defining the width of the terminal, to adjust help formatting to the terminal size of a
|
||||
connected client.
|
||||
Finally, it offers some convenience methods and makes use of custom exceptions.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def help_formatter_factory(terminal_width: int, base_cls: FmtCls = None) -> FmtCls:
|
||||
"""
|
||||
Constructs and returns a subclass of `argparse.HelpFormatter` with a fixed terminal width argument.
|
||||
|
||||
Although a custom formatter class can be explicitly passed into the `ArgumentParser` constructor, this is not
|
||||
as convenient, when making use of sub-parsers.
|
||||
|
||||
Args:
|
||||
terminal_width:
|
||||
The number of columns of the terminal to which to adjust help formatting.
|
||||
base_cls (optional):
|
||||
The base class to use for inheritance. By default `argparse.ArgumentDefaultsHelpFormatter` is used.
|
||||
|
||||
Returns:
|
||||
The subclass of `base_cls` which fixes the constructor's `width` keyword-argument to `terminal_width`.
|
||||
"""
|
||||
if base_cls is None:
|
||||
base_cls = ArgumentDefaultsHelpFormatter
|
||||
|
||||
class ClientHelpFormatter(base_cls):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
kwargs['width'] = terminal_width
|
||||
super().__init__(*args, **kwargs)
|
||||
return ClientHelpFormatter
|
||||
|
||||
def __init__(self, parent: 'CommandParser' = None, **kwargs) -> None:
|
||||
"""
|
||||
Sets additional internal attributes depending on whether a parent-parser was defined.
|
||||
|
||||
The `help_formatter_factory` is called and the returned class is mapped to the `FORMATTER_CLASS` keyword.
|
||||
By default, `exit_on_error` is set to `False` (as opposed to how the parent class handles it).
|
||||
|
||||
Args:
|
||||
parent (optional):
|
||||
An instance of the same class. Intended to be passed as a keyword-argument into the `add_parser` method
|
||||
of the subparsers action returned by the `ArgumentParser.add_subparsers` method. If this is present,
|
||||
the `SESSION_WRITER` and `CLIENT_INFO.TERMINAL_WIDTH` keywords must not be present in `kwargs`.
|
||||
**kwargs(optional):
|
||||
In addition to the regular `ArgumentParser` constructor parameters, this method expects the instance of
|
||||
the `StreamWriter` as well as the terminal width both to be passed explicitly, if the `parent` argument
|
||||
is empty.
|
||||
"""
|
||||
self._session_writer: StreamWriter = parent.session_writer if parent else kwargs.pop(SESSION_WRITER)
|
||||
self._terminal_width: int = parent.terminal_width if parent else kwargs.pop(CLIENT_INFO.TERMINAL_WIDTH)
|
||||
kwargs[FORMATTER_CLASS] = self.help_formatter_factory(self._terminal_width, kwargs.get(FORMATTER_CLASS))
|
||||
kwargs.setdefault('exit_on_error', False)
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@property
|
||||
def session_writer(self) -> StreamWriter:
|
||||
"""Returns the predefined stream writer object of the control session."""
|
||||
return self._session_writer
|
||||
|
||||
@property
|
||||
def terminal_width(self) -> int:
|
||||
"""Returns the predefined terminal width."""
|
||||
return self._terminal_width
|
||||
|
||||
def _print_message(self, message: str, *args, **kwargs) -> None:
|
||||
"""This is overridden to ensure that no messages are sent to stdout/stderr, but always to the stream writer."""
|
||||
if message:
|
||||
self._session_writer.write(message.encode())
|
||||
|
||||
def exit(self, status: int = 0, message: str = None) -> None:
|
||||
"""This is overridden to prevent system exit to be invoked."""
|
||||
if message:
|
||||
self._print_message(message)
|
||||
|
||||
def print_help(self, file=None) -> None:
|
||||
"""This just adds the custom `HelpRequested` exception after the parent class' method."""
|
||||
super().print_help(file)
|
||||
raise HelpRequested
|
||||
|
||||
def add_optional_num_argument(self, *name_or_flags: str, **kwargs) -> Action:
|
||||
"""Convenience method for `add_argument` setting the name, `nargs`, `default`, and `type`, unless specified."""
|
||||
if not name_or_flags:
|
||||
name_or_flags = (NUM, )
|
||||
kwargs.setdefault('nargs', '?')
|
||||
kwargs.setdefault('default', 1)
|
||||
kwargs.setdefault('type', int)
|
||||
return self.add_argument(*name_or_flags, **kwargs)
|
||||
Reference in New Issue
Block a user