diff --git a/pyproject.toml b/pyproject.toml index aed0312..f893c9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,9 +3,10 @@ name = "xts" version = "2.0.0" dependencies = [ "rich>=13.7.1", - "yaml_runner @ git+https://github.com/rdkcentral/yaml_runner.git@2.0.0", + "yaml_runner @ git+https://github.com/rdkcentral/yaml_runner.git@2.1.0", "pyinstaller>=6.12.0", - "requests>=2.32.3" + "requests>=2.32.3", + "rich-argparse>=1.7.2" ] [tool.setuptools.packages.find] include = ["xts_core", "xts_core.plugins"] @@ -14,4 +15,4 @@ namespaces = false [project.scripts] xts-install = "xts_core.install:main" -xts = "xts_core.xts:main" \ No newline at end of file +xts = "xts_core.xts:main" diff --git a/src/xts_core/plugins/base_plugin.py b/src/xts_core/plugins/base_plugin.py index 7b2b252..53abb28 100644 --- a/src/xts_core/plugins/base_plugin.py +++ b/src/xts_core/plugins/base_plugin.py @@ -11,11 +11,16 @@ except: import utils as plugin_utils +try: + from xts_core.xts_arg_parser import XTSArgumentParser +except ImportError: + from xts_arg_parser import XTSArgumentParser + class BaseXTSPlugin(ABC): config_dir = pathlib.Path.home().joinpath('.xts') def __init__(self): - self._prog_parser = argparse.ArgumentParser(prog='xts') + self._prog_parser = XTSArgumentParser(prog='xts') self._subparsers = self._prog_parser.add_subparsers(dest='command') @abstractmethod diff --git a/src/xts_core/plugins/xts_allocator_client.py b/src/xts_core/plugins/xts_allocator_client.py index 2234048..41ebd32 100755 --- a/src/xts_core/plugins/xts_allocator_client.py +++ b/src/xts_core/plugins/xts_allocator_client.py @@ -9,9 +9,9 @@ import rich from rich.table import Table try: - from base_plugin import BaseXTSPlugin, plugin_utils + from base_plugin import BaseXTSPlugin, plugin_utils, XTSArgumentParser except: - from xts_core.plugins.base_plugin import BaseXTSPlugin, plugin_utils + from xts_core.plugins.base_plugin import BaseXTSPlugin, plugin_utils, XTSArgumentParser @@ -163,7 +163,7 @@ def _setup_allocator_args(self): # remove # Since a lot of these parsers share arguments we can use the base parsers as parents # to save on adding the same arguments multiple times. - remove_base = argparse.ArgumentParser('remove base', add_help=False,) + remove_base = XTSArgumentParser('remove base', add_help=False,) remove_base.add_argument('name', help='Name of the server to remove') remove_parser = allocator_subparsers.add_parser('remove', help='Remove an allocator server', @@ -176,10 +176,10 @@ def _setup_allocator_args(self): add_parser.add_argument('url', help='URL for the allocator server') # list allocator_subparsers.add_parser('list', help='List all known allocators') - base_parser = argparse.ArgumentParser('base', add_help=False) + base_parser = XTSArgumentParser('base', add_help=False) base_parser.add_argument('--server', required=True, help='Server URL') # search - search_base = argparse.ArgumentParser('search_base', + search_base = XTSArgumentParser('search_base', add_help=False, parents=[base_parser]) search_base.add_argument('--platform', help='Platform of DUT') @@ -189,7 +189,7 @@ def _setup_allocator_args(self): help='Search for slots on an allocator server', parents=[search_base]) # add-slot - add_base = argparse.ArgumentParser('add_base', + add_base = XTSArgumentParser('add_base', add_help=False, parents=[search_base]) add_base.add_argument('--owner-email', help='Owner email (if allocated).', dest='owner_email') diff --git a/src/xts_core/xts.py b/src/xts_core/xts.py index a702322..5af47d8 100755 --- a/src/xts_core/xts.py +++ b/src/xts_core/xts.py @@ -33,7 +33,7 @@ The script utilizes the `yaml_runner` module to handle configuration parsing and command execution. """ -import argparse + import os import re import sys @@ -55,14 +55,19 @@ try: from .utils import info, error, warning, is_url -except: +except ImportError: from xts_core.utils import info, error, warning, is_url try: from . import xts_alias -except: +except ImportError: from xts_core import xts_alias +try: + from .xts_arg_parser import XTSArgumentParser +except ImportError: + from xts_core.xts_arg_parser import XTSArgumentParser + class XTS(): """ @@ -136,7 +141,7 @@ def _parse_first_arg(self): list: Remaining arguments after parsing, including the command name. """ # quick parser for alias commands - pre_parser = argparse.ArgumentParser(prog="xts", add_help=True) + pre_parser = XTSArgumentParser(prog="xts", add_help=True) pre_subparsers = pre_parser.add_subparsers(dest="command", required=True) self._add_alias_subcommands(pre_subparsers) @@ -155,7 +160,7 @@ def _parse_first_arg(self): self._find_xts_config() # full parser with YAML/plugin commands - parser = argparse.ArgumentParser(prog="xts") + parser = XTSArgumentParser(prog="xts") subparsers = parser.add_subparsers(dest="command", required=True) self._add_alias_subcommands(subparsers) @@ -379,7 +384,8 @@ def run(self): yaml_runner = YamlRunner(self._command_sections, program='xts', hierarchical=True, - fail_fast=True) + fail_fast=True, + parser_class=XTSArgumentParser) _,_,exit_code = yaml_runner.run(args) sys.exit(sorted(exit_code)[-1]) except Exception as e: diff --git a/src/xts_core/xts_arg_parser.py b/src/xts_core/xts_arg_parser.py new file mode 100644 index 0000000..6358cda --- /dev/null +++ b/src/xts_core/xts_arg_parser.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 +#** ***************************************************************************** +# * +# * If not stated otherwise in this file or this component's LICENSE file the +# * following copyright and licenses apply: +# * +# * Copyright 2024 RDK Management +# * +# * Licensed under the Apache License, Version 2.0 (the "License"); +# * you may not use this file except in compliance with the License. +# * You may obtain a copy of the License at +# * +# * +# http://www.apache.org/licenses/LICENSE-2.0 +# * +# * Unless required by applicable law or agreed to in writing, software +# * distributed under the License is distributed on an "AS IS" BASIS, +# * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# * See the License for the specific language governing permissions and +# * limitations under the License. +# * +#* ****************************************************************************** + +import argparse +import sys + +from rich_argparse import RichHelpFormatter + +class XTSArgumentParser(argparse.ArgumentParser): + """XTS Specific argument parser. + """ + + def __init__(self,*args, **kwargs): + kwargs.pop('formatter_class','') + super().__init__(*args, + **kwargs, + formatter_class=RichHelpFormatter) + + def error(self, message): + sys.stderr.write('error: %s\n' % message) + self.print_help() + sys.exit(2) \ No newline at end of file