223 lines
8.2 KiB
Python
223 lines
8.2 KiB
Python
|
# 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.
|
||
|
|
||
|
"""
|
||
|
This module provides a drop-in replacement for the standard library argparse
|
||
|
module, with the output of the 'help' action automatically sent to a pager
|
||
|
when appropriate.
|
||
|
|
||
|
To use, replace the code:
|
||
|
|
||
|
>>> import argparse
|
||
|
|
||
|
with:
|
||
|
|
||
|
>>> from autopage import argparse
|
||
|
|
||
|
Or, alternatively, call the ``autopage.argparse.monkey_patch()`` function to
|
||
|
monkey-patch the argparse module. This is useful when you do not control the
|
||
|
code that creates the ArgumentParser. The result of calling this function can
|
||
|
also be used as a context manager to ensure that the original functionality is
|
||
|
restored.
|
||
|
"""
|
||
|
|
||
|
import argparse
|
||
|
import contextlib
|
||
|
import functools
|
||
|
import types
|
||
|
from typing import Any, Sequence, Text, TextIO, Tuple, Type, Optional, Union
|
||
|
from typing import Callable, ContextManager, Generator
|
||
|
|
||
|
import autopage
|
||
|
|
||
|
from argparse import * # noqa
|
||
|
|
||
|
|
||
|
_HelpFormatter = argparse.HelpFormatter
|
||
|
|
||
|
_color_attr = '_autopage_color'
|
||
|
|
||
|
|
||
|
def help_pager(out_stream: Optional[TextIO] = None) -> autopage.AutoPager:
|
||
|
"""Return an AutoPager suitable for help output."""
|
||
|
return autopage.AutoPager(out_stream,
|
||
|
allow_color=True,
|
||
|
line_buffering=False,
|
||
|
reset_on_exit=False)
|
||
|
|
||
|
|
||
|
def use_color_for_parser(parser: argparse.ArgumentParser,
|
||
|
color: bool) -> None:
|
||
|
"""Configure a parser whether to output in color from HelpFormatters."""
|
||
|
setattr(parser, _color_attr, color)
|
||
|
|
||
|
|
||
|
class ColorHelpFormatter(_HelpFormatter):
|
||
|
class _Section(_HelpFormatter._Section): # type: ignore
|
||
|
@property
|
||
|
def heading(self) -> Optional[Text]:
|
||
|
if (not self._heading
|
||
|
or self._heading == argparse.SUPPRESS
|
||
|
or not getattr(self.formatter, _color_attr, False)):
|
||
|
return self._heading
|
||
|
return f'\033[4m{self._heading}\033[0m'
|
||
|
|
||
|
@heading.setter
|
||
|
def heading(self, heading: Optional[Text]) -> None:
|
||
|
self._heading = heading
|
||
|
|
||
|
def _metavar_formatter(self,
|
||
|
action: argparse.Action,
|
||
|
default_metavar: Text) -> Callable[[int],
|
||
|
Tuple[str, ...]]:
|
||
|
get_metavars = super()._metavar_formatter(action, default_metavar)
|
||
|
if not getattr(self, _color_attr, False):
|
||
|
return get_metavars
|
||
|
|
||
|
def color_metavar(size: int) -> Tuple[str, ...]:
|
||
|
return tuple(f'\033[3m{mv}\033[0m' for mv in get_metavars(size))
|
||
|
return color_metavar
|
||
|
|
||
|
|
||
|
class ColorRawDescriptionHelpFormatter(ColorHelpFormatter,
|
||
|
argparse.RawDescriptionHelpFormatter):
|
||
|
"""Help message formatter which retains any formatting in descriptions."""
|
||
|
|
||
|
|
||
|
class ColorRawTextHelpFormatter(ColorHelpFormatter,
|
||
|
argparse.RawTextHelpFormatter):
|
||
|
"""Help message formatter which retains formatting of all help text."""
|
||
|
|
||
|
|
||
|
class ColorArgDefaultsHelpFormatter(ColorHelpFormatter,
|
||
|
argparse.ArgumentDefaultsHelpFormatter):
|
||
|
"""Help message formatter which adds default values to argument help."""
|
||
|
|
||
|
|
||
|
class ColorMetavarTypeHelpFormatter(ColorHelpFormatter,
|
||
|
argparse.MetavarTypeHelpFormatter):
|
||
|
"""Help message formatter which uses the argument 'type' as the default
|
||
|
metavar value (instead of the argument 'dest')"""
|
||
|
|
||
|
|
||
|
class _HelpAction(argparse._HelpAction):
|
||
|
def __init__(self,
|
||
|
option_strings: Sequence[Text],
|
||
|
dest: Text = argparse.SUPPRESS,
|
||
|
default: Text = argparse.SUPPRESS,
|
||
|
help: Optional[Text] = None) -> None:
|
||
|
argparse.Action.__init__(
|
||
|
self,
|
||
|
option_strings=option_strings,
|
||
|
dest=dest,
|
||
|
default=default,
|
||
|
nargs=0,
|
||
|
help=help)
|
||
|
|
||
|
def __call__(self, parser: argparse.ArgumentParser,
|
||
|
namespace: argparse.Namespace,
|
||
|
values: Union[Text, Sequence[Any], None],
|
||
|
option_string: Optional[Text] = None) -> None:
|
||
|
pager = help_pager()
|
||
|
with pager as out:
|
||
|
use_color_for_parser(parser, pager.to_terminal())
|
||
|
parser.print_help(out)
|
||
|
parser.exit(pager.exit_code())
|
||
|
|
||
|
|
||
|
class _ActionsContainer(argparse._ActionsContainer):
|
||
|
def __init__(self, *args, **kwargs) -> None: # type: ignore
|
||
|
super().__init__(*args, **kwargs)
|
||
|
self.register('action', 'help', _HelpAction)
|
||
|
|
||
|
|
||
|
def _substitute_formatter(
|
||
|
get_fmtr: Callable[[Any], _HelpFormatter]
|
||
|
) -> Callable[[argparse.ArgumentParser], _HelpFormatter]:
|
||
|
@functools.wraps(get_fmtr)
|
||
|
def _get_formatter(parser: argparse.ArgumentParser) -> _HelpFormatter:
|
||
|
if parser.formatter_class is _HelpFormatter:
|
||
|
parser.formatter_class = ColorHelpFormatter
|
||
|
formatter = get_fmtr(parser)
|
||
|
if isinstance(formatter, ColorHelpFormatter):
|
||
|
setattr(formatter, _color_attr,
|
||
|
getattr(parser, _color_attr, False))
|
||
|
return formatter
|
||
|
return _get_formatter
|
||
|
|
||
|
|
||
|
class AutoPageArgumentParser(argparse.ArgumentParser, _ActionsContainer):
|
||
|
@_substitute_formatter
|
||
|
def _get_formatter(self) -> _HelpFormatter:
|
||
|
return super()._get_formatter()
|
||
|
|
||
|
|
||
|
ArgumentParser = AutoPageArgumentParser # type: ignore
|
||
|
HelpFormatter = ColorHelpFormatter # type: ignore
|
||
|
RawDescriptionHelpFormatter = ColorRawDescriptionHelpFormatter # type: ignore
|
||
|
RawTextHelpFormatter = ColorRawTextHelpFormatter # type: ignore
|
||
|
ArgumentDefaultsHelpFormatter = ColorArgDefaultsHelpFormatter # type: ignore
|
||
|
MetavarTypeHelpFormatter = ColorMetavarTypeHelpFormatter # type: ignore
|
||
|
|
||
|
|
||
|
def monkey_patch() -> ContextManager:
|
||
|
"""
|
||
|
Monkey-patch the system argparse module to automatically page help output.
|
||
|
|
||
|
The result of calling this function can optionally be used as a context
|
||
|
manager to restore the status quo when it exits.
|
||
|
"""
|
||
|
import sys
|
||
|
|
||
|
def get_existing_classes(module: types.ModuleType) -> Tuple[Type, ...]:
|
||
|
return (
|
||
|
module._HelpAction, # type: ignore
|
||
|
module.HelpFormatter, # type: ignore
|
||
|
module.RawDescriptionHelpFormatter, # type: ignore
|
||
|
module.RawTextHelpFormatter, # type: ignore
|
||
|
module.ArgumentDefaultsHelpFormatter, # type: ignore
|
||
|
module.MetavarTypeHelpFormatter, # type: ignore
|
||
|
) # type: ignore
|
||
|
|
||
|
def patch_classes(module: types.ModuleType,
|
||
|
impl: Tuple[Type, ...]) -> None:
|
||
|
(
|
||
|
module._HelpAction, # type: ignore
|
||
|
module.HelpFormatter, # type: ignore
|
||
|
module.RawDescriptionHelpFormatter, # type: ignore
|
||
|
module.RawTextHelpFormatter, # type: ignore
|
||
|
module.ArgumentDefaultsHelpFormatter, # type: ignore
|
||
|
module.MetavarTypeHelpFormatter, # type: ignore
|
||
|
) = impl
|
||
|
|
||
|
orig = get_existing_classes(argparse)
|
||
|
orig_fmtr = argparse.ArgumentParser._get_formatter
|
||
|
patched = get_existing_classes(sys.modules[__name__])
|
||
|
patch_classes(argparse, patched)
|
||
|
new_fmtr = _substitute_formatter(orig_fmtr)
|
||
|
argparse.ArgumentParser._get_formatter = new_fmtr # type: ignore
|
||
|
|
||
|
@contextlib.contextmanager
|
||
|
def unpatcher() -> Generator:
|
||
|
try:
|
||
|
yield
|
||
|
finally:
|
||
|
patch_classes(argparse, orig)
|
||
|
argparse.ArgumentParser._get_formatter = orig_fmtr # type: ignore
|
||
|
|
||
|
return unpatcher()
|
||
|
|
||
|
|
||
|
__all__ = argparse.__all__ + [ # type: ignore
|
||
|
'use_color_for_parser', 'monkey_patch'
|
||
|
]
|